From 1e0c8b39548953f144c399d29b7b3e568ca8dc11 Mon Sep 17 00:00:00 2001 From: tylerbishopdev <78498397+tylerbishopdev@users.noreply.github.com> Date: Sun, 8 Nov 2026 19:48:46 -0800 Subject: [PATCH] fix(agent-plugins): scope tool discovery to requested servers --- packages/bb-plugin-agent-plugins/README.md | 2 + packages/bb-plugin-agent-plugins/server.ts | 8 +- .../src/bridge-scoped.test.ts | 72 +++++ .../src/gateway-scoped.test.ts | 301 ++++++++++++++++++ .../bb-plugin-agent-plugins/src/gateway.ts | 74 ++++- 5 files changed, 437 insertions(+), 20 deletions(-) create mode 100644 packages/bb-plugin-agent-plugins/src/bridge-scoped.test.ts create mode 100644 packages/bb-plugin-agent-plugins/src/gateway-scoped.test.ts diff --git a/packages/bb-plugin-agent-plugins/README.md b/packages/bb-plugin-agent-plugins/README.md index 650151e..8e8a42b 100644 --- a/packages/bb-plugin-agent-plugins/README.md +++ b/packages/bb-plugin-agent-plugins/README.md @@ -36,6 +36,8 @@ bb plugin install git:https://github.com/patleeman/bb-plugins.git --subdirectory - **CLI:** `bb agent-plugins list --json` · `bb agent-plugins outdated --json` · `bb agent-plugins update ` · `bb agent-plugins show --json` · `bb agent-plugins tools --json` · `bb agent-plugins tools --query "send message" --server slack` · `bb agent-plugins describe ` · `bb agent-plugins call '{"text":"hello"}'` - **Agents:** the bridge exposes `agent_plugins_list_tools`, `agent_plugins_describe_tool`, and `agent_plugins_call`. `agent_plugins_list_tools` searches the catalog (`query`, `serverId`, `limit`, `offset`) and returns compact entries without input schemas, plus a tool count per server, so large servers don't flood the context. Agents fetch schemas for the tools they pick with `agent_plugins_describe_tool`, then call them by opaque ID. +Tool discovery applies `serverId` before loading catalogs; scoped results include counts for that scope only. Unscoped discovery loads at most four server catalogs concurrently per request and preserves catalog order. Describe/call resolve IDs against their owning server, reusing the existing per-server cache, refresh notifications, five-minute TTL, and five-second error backoff. Unrelated failed servers are not retried by targeted requests. After a gateway restart, an ID with an ambiguous readable server prefix requires scoped discovery again; full IDs are always checked against the current catalog before execution. + The MCP client delegates authentication to the official SDK, including protected-resource/authorization-server discovery, dynamic client registration, PKCE, refresh, `WWW-Authenticate`, insufficient-scope step-up, and OAuth callback completion. Tokens and PKCE state are stored in a BB secret setting; they are never placed in the database snapshot or frontend payload. For OAuth-protected MCP servers such as Fastmail, the registered redirect uses `http://localhost` on BB's port. If BB runs on another machine, sign in with your normal browser, then copy the **full** localhost URL from its address bar after the connection fails and paste it into the matching server's **Finish sign-in** field in BB. The URL contains a short-lived, one-use authorization code: paste it only into BB, never into chat or logs. BB verifies the callback URL, state and PKCE grant before exchanging the code. If the browser runs on the BB host or you have a port forward, the callback completes normally without pasting. Fastmail's MCP dynamic registration does not accept arbitrary public HTTPS redirect URIs; setting `BB_APP_URL` does not change this callback. diff --git a/packages/bb-plugin-agent-plugins/server.ts b/packages/bb-plugin-agent-plugins/server.ts index 706d462..6b53a4c 100644 --- a/packages/bb-plugin-agent-plugins/server.ts +++ b/packages/bb-plugin-agent-plugins/server.ts @@ -559,7 +559,7 @@ export default async function plugin(bb: BbPluginApi) { limit: z.number().int().min(1).max(MAX_TOOL_SEARCH_LIMIT).optional().describe(`Max tools to return (default ${DEFAULT_TOOL_SEARCH_LIMIT}).`), offset: z.number().int().min(0).optional().describe("Skip this many matches, for paging."), }).strict(), - async execute(input) { return JSON.stringify(searchTools(await gateway.listTools(), input), null, 2); }, + async execute(input) { return JSON.stringify(searchTools(await gateway.listTools(input.serverId), input), null, 2); }, }); bb.agents.registerTool({ name: "agent_plugins_describe_tool", @@ -567,7 +567,7 @@ export default async function plugin(bb: BbPluginApi) { instructions: "Pass opaqueIds returned by agent_plugins_list_tools; read inputSchema before calling agent_plugins_call.", presentation: { label: { pending: "Describing Agent Plugins tools", completed: "Tools described" } }, parameters: z.object({ opaqueIds: z.array(z.string().min(1)).min(1).max(MAX_DESCRIBE_TOOLS) }).strict(), - async execute(input) { return JSON.stringify(describeTools(await gateway.listTools(), input.opaqueIds), null, 2); }, + async execute(input) { return JSON.stringify(describeTools(await gateway.resolveTools(input.opaqueIds), input.opaqueIds), null, 2); }, }); bb.agents.registerTool({ name: "agent_plugins_call", @@ -1270,7 +1270,7 @@ export default async function plugin(bb: BbPluginApi) { }; const query = flag("--query"); const serverId = flag("--server"); const limit = flag("--limit"); const offset = flag("--offset"); - const tools = await gateway.listTools(); + const tools = await gateway.listTools(serverId); if (query === undefined && serverId === undefined && limit === undefined && offset === undefined) { return { exitCode: 0, stdout: (asJson ? JSON.stringify({ tools }, null, 2) : tools.map((t: CatalogTool) => `${t.opaqueId} — ${t.description} [${t.serverType}]`).join("\n")) + "\n" }; } @@ -1288,7 +1288,7 @@ export default async function plugin(bb: BbPluginApi) { if (cmd === "describe") { const ids = rest.filter((arg) => arg !== "--json"); if (ids.length === 0) return { exitCode: 2, stderr: "Usage: bb agent-plugins describe … [--json]\n" }; - const result = describeTools(await gateway.listTools(), ids); + const result = describeTools(await gateway.resolveTools(ids), ids); const text = [ ...result.tools.map((t) => `${t.opaqueId}\n${t.description}\ninputSchema: ${JSON.stringify(t.inputSchema)}`), ...result.notFound.map((id) => `Not found: ${id}`), diff --git a/packages/bb-plugin-agent-plugins/src/bridge-scoped.test.ts b/packages/bb-plugin-agent-plugins/src/bridge-scoped.test.ts new file mode 100644 index 0000000..30c980d --- /dev/null +++ b/packages/bb-plugin-agent-plugins/src/bridge-scoped.test.ts @@ -0,0 +1,72 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createFakePluginHost } from "@get-bb/plugin-sdk/testing"; +import plugin from "../server.js"; +import { McpGateway } from "./gateway.js"; +import type { CatalogTool } from "./types.js"; + +const hosts: Array> = []; +const tool: CatalogTool = { + opaqueId: "fixture__healthy__echo_0123456789", pluginId: "fixture", pluginName: "Fixture", + serverId: "healthy", serverType: "stdio", name: "echo", description: "Echo a value.", + inputSchema: { type: "object", properties: { value: { type: "string" } } }, status: "ready", +}; + +async function createHost() { + const host = createFakePluginHost({ pluginId: "agent-plugins" }); + hosts.push(host); + await plugin(host.bb); + return host.harness.behavior; +} + +afterEach(async () => { + await Promise.all(hosts.splice(0).map((host) => host.harness.lifecycle.dispose())); + vi.restoreAllMocks(); +}); + +describe("scoped static bridge contracts", () => { + it("passes the server filter to discovery before searching; CLI uses the same scope", async () => { + const list = vi.spyOn(McpGateway.prototype, "listTools").mockResolvedValue([tool]); + const host = await createHost(); + const result = JSON.parse(String(await host.callAgentTool("agent_plugins_list_tools", { serverId: "healthy", query: "echo", limit: 1, offset: 0 }))); + expect(list).toHaveBeenCalledExactlyOnceWith("healthy"); + expect(result).toMatchObject({ total: 1, offset: 0, tools: [{ opaqueId: tool.opaqueId }], servers: [{ pluginName: "Fixture", serverId: "healthy", tools: 1 }] }); + expect(result.tools[0]).not.toHaveProperty("inputSchema"); + list.mockClear(); + const cli = await host.runCli(["tools", "--server", "healthy", "--json"]); + expect(cli.exitCode).toBe(0); + expect(list).toHaveBeenCalledExactlyOnceWith("healthy"); + expect(JSON.parse(cli.stdout!)).toEqual(result); + }); + + it("describes requested IDs without invoking global discovery, preserving notFound and deduplication", async () => { + const list = vi.spyOn(McpGateway.prototype, "listTools"); + const resolve = vi.spyOn(McpGateway.prototype, "resolveTools").mockResolvedValue([tool]); + const host = await createHost(); + const ids = [tool.opaqueId, "missing", tool.opaqueId]; + const result = JSON.parse(String(await host.callAgentTool("agent_plugins_describe_tool", { opaqueIds: ids }))); + expect(resolve).toHaveBeenCalledExactlyOnceWith(ids); + expect(list).not.toHaveBeenCalled(); + expect(result).toEqual({ tools: [tool], notFound: ["missing"] }); + resolve.mockClear(); + const cli = await host.runCli(["describe", ...ids, "--json"]); + expect(cli.exitCode).toBe(1); + expect(resolve).toHaveBeenCalledExactlyOnceWith(ids); + expect(JSON.parse(cli.stdout!)).toEqual(result); + expect(list).not.toHaveBeenCalled(); + }); + + it("retains strict input schemas, call defaults, cancellation and the error envelope", async () => { + const list = vi.spyOn(McpGateway.prototype, "listTools"); + const call = vi.spyOn(McpGateway.prototype, "call").mockRejectedValue(new Error("fixture call error")); + const host = await createHost(); + await expect(host.callAgentTool("agent_plugins_list_tools", { serverId: "healthy", unknown: true })).rejects.toThrow(); + await expect(host.callAgentTool("agent_plugins_list_tools", { limit: 201 })).rejects.toThrow(); + await expect(host.callAgentTool("agent_plugins_describe_tool", { opaqueIds: [] })).rejects.toThrow(); + await expect(host.callAgentTool("agent_plugins_describe_tool", { opaqueIds: Array(21).fill("id") })).rejects.toThrow(); + expect(list).not.toHaveBeenCalled(); + const signal = new AbortController().signal; + const result = JSON.parse(String(await host.callAgentTool("agent_plugins_call", { opaqueId: tool.opaqueId }, { signal }))); + expect(call).toHaveBeenCalledExactlyOnceWith(tool.opaqueId, {}, signal); + expect(result).toEqual({ isError: true, error: "fixture call error" }); + }); +}); diff --git a/packages/bb-plugin-agent-plugins/src/gateway-scoped.test.ts b/packages/bb-plugin-agent-plugins/src/gateway-scoped.test.ts new file mode 100644 index 0000000..ca07ab5 --- /dev/null +++ b/packages/bb-plugin-agent-plugins/src/gateway-scoped.test.ts @@ -0,0 +1,301 @@ +import Database from "better-sqlite3"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { McpGateway, type McpStdioCatalog, type McpStdioHost } from "./gateway.js"; +import { AgentPluginsStore } from "./store.js"; +import { describeTools, searchTools } from "./tool-search.js"; +import type { PluginRecord } from "./types.js"; + +const databases: Database.Database[] = []; +const gateways: McpGateway[] = []; +const log = { info() {}, warn() {}, error() {} }; + +function catalog(name = "echo"): McpStdioCatalog { + return { + tools: [{ + name, + description: "Echo a value. More detail.", + inputSchema: { type: "object", properties: { value: { type: "string" } }, required: ["value"] }, + outputSchema: { type: "object", properties: { value: { type: "string" } } }, + annotations: { readOnlyHint: true }, + _meta: { fixture: true }, + }], + prompts: [], resources: [], resourceTemplates: [], + }; +} + +function fixture() { + const database = new Database(":memory:"); + databases.push(database); + const store = new AgentPluginsStore(database, (db, statements) => { + for (const statement of statements) db.exec(statement); + }); + const plugin: PluginRecord = { + id: "fixture-plugin", name: "Fixture", version: "1.0.0", description: null, + specVersion: "1.0.0", sourceType: "path", sourceIntent: "path:/fixture", + sourceResolved: null, sourceRef: null, tagPrefix: null, + pluginRoot: "/fixture", pluginData: "/fixture/data", activeGen: 1, + status: "active", approval: "approved", lastError: null, contentHash: null, + installedAt: 1, updatedAt: 1, + }; + store.upsertPlugin(plugin); + const add = (serverId: string, pluginId = plugin.id) => { + if (pluginId !== plugin.id) store.upsertPlugin({ ...plugin, id: pluginId, name: pluginId }); + const record = { + pluginId, serverId, type: "stdio" as const, + configJson: JSON.stringify({ type: "stdio", command: "fixture-only-never-spawned" }), + status: "idle" as const, lastError: null, approved: 1, enabled: 1, + }; + store.upsertMcpServer(record); + return record; + }; + const catalogs = new Map(); + const failures = new Set(); + const stalls = new Set(); + const current = (key: string) => catalogs.get(key) ?? catalog(); + let changed: Parameters>[0] | undefined; + const host = { + start: vi.fn(async ({ key }, signal) => { + if (failures.has(key)) throw new Error(`fixture failure: ${key}`); + if (stalls.has(key)) return new Promise((_resolve, reject) => { + if (signal?.aborted) reject(signal.reason); + else signal?.addEventListener("abort", () => reject(signal.reason), { once: true }); + }); + return current(key); + }), + refresh: vi.fn(async (key) => current(key)), + close: vi.fn(async () => {}), + callTool: vi.fn(async (_key, _name, args) => ({ + content: [{ type: "text", text: "echoed" }], structuredContent: args, _meta: { fixture: true }, + })), + getPrompt: async () => ({}), readResource: async () => ({}), complete: async () => ({}), + subscribeResource: async () => {}, unsubscribeResource: async () => {}, setLoggingLevel: async () => {}, + onCatalogChanged(handler) { changed = handler; return () => { changed = undefined; }; }, + } satisfies McpStdioHost; + const createGateway = () => { + const gateway = new McpGateway(store, log, { stdioHost: host }); + gateways.push(gateway); + return gateway; + }; + const starts = () => host.start.mock.calls.map(([config]) => config.key); + return { store, plugin, add, host, catalogs, failures, stalls, createGateway, starts, + changed: (key: string) => changed?.(key, "tools", null) }; +} + +afterEach(async () => { + await Promise.all(gateways.splice(0).map((gateway) => gateway.close())); + for (const database of databases.splice(0)) database.close(); + vi.useRealTimers(); +}); + +describe("scoped MCP tool discovery", () => { + it.each(["error", "hang"])("cold and warm server scopes never start an unrelated %s server", async (mode) => { + const f = fixture(); + f.add("a-broken"); f.add("healthy"); + (mode === "error" ? f.failures : f.stalls).add("fixture-plugin:a-broken"); + const gateway = f.createGateway(); + expect(await gateway.listTools("missing")).toEqual([]); + const cold = await gateway.listTools(" healthy "); + expect(cold.map((tool) => tool.name)).toEqual(["echo"]); + expect(await gateway.listTools("healthy")).toEqual(cold); + expect(f.starts()).toEqual(["fixture-plugin:healthy"]); + expect(f.host.refresh).not.toHaveBeenCalled(); + expect(searchTools(cold, { serverId: " healthy ", query: "echo", limit: 1 })).toEqual({ + total: 1, offset: 0, + tools: [{ opaqueId: cold[0]!.opaqueId, pluginName: "Fixture", serverId: "healthy", name: "echo", summary: "Echo a value.", status: "ready" }], + servers: [{ pluginName: "Fixture", serverId: "healthy", tools: 1 }], + }); + }); + + it.each(["describe", "call"])("resolves a previous gateway's opaque ID for cold %s without a global walk", async (operation) => { + const f = fixture(); f.add("healthy__server"); + f.catalogs.set("fixture-plugin:healthy__server", catalog("echo__value")); + const old = f.createGateway(); + const [definition] = await old.listTools("healthy__server"); + await old.close(); + f.add("a-broken"); f.stalls.add("fixture-plugin:a-broken"); + f.host.start.mockClear(); + const gateway = f.createGateway(); + if (operation === "describe") { + const ids = [definition!.opaqueId, "unknown", definition!.opaqueId]; + expect(describeTools(await gateway.resolveTools(ids), ids)).toEqual({ tools: [definition], notFound: ["unknown"] }); + } else { + const signal = new AbortController().signal; + expect(await gateway.call(definition!.opaqueId, { value: "hello" }, signal)).toEqual({ + content: [{ type: "text", text: "echoed" }], structuredContent: { value: "hello" }, _meta: { fixture: true }, + }); + expect(f.host.callTool).toHaveBeenCalledExactlyOnceWith("fixture-plugin:healthy__server", "echo__value", { value: "hello" }, catalog("echo__value").tools[0], signal); + } + expect(f.starts()).toEqual(["fixture-plugin:healthy__server"]); + }); + + it("keeps targeted discovery and calls responsive while a broad cold discovery is stalled", async () => { + vi.useFakeTimers(); + const f = fixture(); f.add("a-broken"); f.add("healthy"); + f.stalls.add("fixture-plugin:a-broken"); + const gateway = f.createGateway(); + let finished = false; + const broad = gateway.listTools().then((tools) => { finished = true; return tools; }); + const [healthy] = await gateway.listTools("healthy"); + await gateway.call(healthy!.opaqueId, { value: "hello" }); + expect(finished).toBe(false); + expect(f.starts().filter((key) => key.endsWith(":healthy"))).toHaveLength(1); + await vi.advanceTimersByTimeAsync(15_000); + expect((await broad).map((tool) => [tool.serverId, tool.status])).toEqual([["a-broken", "error"], ["healthy", "ready"]]); + // A timeout is cached per server, even though a targeted healthy call is free to proceed. + await gateway.listTools(); + expect(f.starts().filter((key) => key.endsWith(":a-broken"))).toHaveLength(1); + }); + + it("bounds broad discovery concurrency and preserves catalog order despite out-of-order responses", async () => { + vi.useFakeTimers(); + const f = fixture(); + for (let i = 0; i < 7; i++) f.add(`server-${i}`); + let active = 0; let peak = 0; + f.host.start.mockImplementation(async ({ key }) => { + active++; peak = Math.max(peak, active); + await new Promise((resolve) => setTimeout(resolve, key.endsWith("0") ? 50 : 10)); + active--; + return catalog(key); + }); + const result = f.createGateway().listTools(); + await vi.advanceTimersByTimeAsync(0); + expect(f.host.start).toHaveBeenCalledTimes(4); + await vi.advanceTimersByTimeAsync(50); + expect((await result).map((tool) => tool.serverId)).toEqual(Array.from({ length: 7 }, (_, i) => `server-${i}`)); + expect(peak).toBe(4); + }); + + it("backs off only the failed server; healthy lookups never retry it, even after its backoff expires", async () => { + vi.useFakeTimers(); + const f = fixture(); f.add("a-broken"); f.add("healthy"); + f.failures.add("fixture-plugin:a-broken"); + const gateway = f.createGateway(); + const tools = await gateway.listTools(); + const broken = tools.find((tool) => tool.status === "error")!; + const healthy = tools.find((tool) => tool.status === "ready")!; + expect(describeTools(await gateway.resolveTools([broken.opaqueId]), [broken.opaqueId]).tools).toEqual([broken]); + await expect(gateway.call(broken.opaqueId, {})).rejects.toThrow("fixture failure"); + await gateway.listTools(); + expect(f.starts()).toHaveLength(2); + await vi.advanceTimersByTimeAsync(5_000); + await gateway.resolveTools([healthy.opaqueId]); + await gateway.call(healthy.opaqueId, {}); + expect(f.starts()).toHaveLength(2); + await gateway.listTools("a-broken"); + expect(f.starts().filter((key) => key.endsWith(":a-broken"))).toHaveLength(2); + expect(f.host.callTool).toHaveBeenCalledTimes(1); + }); + + it("refreshes only the owning stale catalog, rejects removed tools, and honors list-changed updates", async () => { + vi.useFakeTimers(); + const f = fixture(); f.add("healthy"); f.add("other"); + const gateway = f.createGateway(); + const tools = await gateway.listTools(); + const old = tools.find((tool) => tool.serverId === "healthy")!; + f.catalogs.set("fixture-plugin:healthy", catalog("replacement")); + await vi.advanceTimersByTimeAsync(5 * 60_000); + expect(await gateway.resolveTools([old.opaqueId])).toEqual([]); + await expect(gateway.call(old.opaqueId, {})).rejects.toThrow("Tool not found"); + expect(f.host.refresh).toHaveBeenCalledExactlyOnceWith("fixture-plugin:healthy", expect.any(AbortSignal)); + const [replacement] = await gateway.listTools("healthy"); + f.catalogs.set("fixture-plugin:healthy", catalog("notified")); + await f.changed("fixture-plugin:healthy"); + await vi.advanceTimersByTimeAsync(0); + expect(await gateway.resolveTools([replacement!.opaqueId])).toEqual([]); + expect((await gateway.listTools("healthy"))[0]!.name).toBe("notified"); + expect(f.host.callTool).not.toHaveBeenCalled(); + }); + + it("rejects unknown, forged, wrong-owner and wrong-kind IDs without executing any tool", async () => { + const f = fixture(); f.add("healthy"); f.add("other"); + f.catalogs.set("fixture-plugin:healthy", { ...catalog(), prompts: [{ name: "echo" }] }); + const gateway = f.createGateway(); + expect(await gateway.resolveTools(["unknown", "fixture-plugin__healthy__echo"])).toEqual([]); + expect(f.host.start).not.toHaveBeenCalled(); + const [tool] = await gateway.listTools("healthy"); + const forged = tool!.opaqueId.slice(0, -10) + "0000000000"; + const wrongOwner = tool!.opaqueId.replace("__healthy__", "__other__"); + // Prompt/resource IDs have a different hash; identical readable names confer no authority. + const wrongKind = (await gateway.listPrompts())[0]!.opaqueId; + for (const id of [forged, wrongOwner, wrongKind]) { + await expect(gateway.call(id, {})).rejects.toThrow("Tool not found"); + } + expect(f.host.callTool).not.toHaveBeenCalled(); + }); + + it.each(["disabled", "unapproved", "removed"])("rejects a known ID after its owner is %s", async (state) => { + const f = fixture(); const record = f.add("healthy"); + const gateway = f.createGateway(); + const [tool] = await gateway.listTools("healthy"); + if (state === "removed") f.store.deletePlugin(f.plugin.id); + else f.store.upsertMcpServer({ ...record, enabled: state === "disabled" ? 0 : 1, approved: state === "unapproved" ? 0 : 1 }); + expect(await gateway.resolveTools([tool!.opaqueId])).toEqual([]); + await expect(gateway.call(tool!.opaqueId, {})).rejects.toThrow("Tool not found"); + expect(f.host.callTool).not.toHaveBeenCalled(); + expect(f.host.start).toHaveBeenCalledTimes(1); + }); + + it("does not guess between colliding slug prefixes; scoped discovery supplies exact ownership", async () => { + const f = fixture(); f.add("healthy.one"); + const old = f.createGateway(); + const [tool] = await old.listTools("healthy.one"); + await old.close(); + f.add("healthy_one"); f.stalls.add("fixture-plugin:healthy_one"); + f.host.start.mockClear(); + const gateway = f.createGateway(); + expect(await gateway.resolveTools([tool!.opaqueId])).toEqual([]); + expect(f.host.start).not.toHaveBeenCalled(); + expect(await gateway.listTools("healthy.one")).toEqual([tool]); + await gateway.call(tool!.opaqueId, {}); + expect(f.starts()).toEqual(["fixture-plugin:healthy.one"]); + }); + + it("revalidates known IDs after an explicit server reset", async () => { + const f = fixture(); f.add("healthy"); + const gateway = f.createGateway(); + const [old] = await gateway.listTools("healthy"); + await gateway.resetServer(f.plugin.id, "healthy"); + f.catalogs.set("fixture-plugin:healthy", catalog("replacement")); + expect(await gateway.resolveTools([old!.opaqueId])).toEqual([]); + expect(f.host.start).toHaveBeenCalledTimes(2); + expect(f.host.callTool).not.toHaveBeenCalled(); + }); + + it("rechecks approval after discovery before using an already-connected server", async () => { + const f = fixture(); const record = f.add("healthy"); + const gateway = f.createGateway(); + const [tool] = await gateway.listTools("healthy"); + await gateway.closeServer(f.plugin.id, "healthy"); + f.host.start.mockImplementation(async () => { + f.store.upsertMcpServer({ ...record, approved: 0 }); + return catalog(); + }); + await expect(gateway.call(tool!.opaqueId, {})).rejects.toThrow("not approved"); + expect(f.host.callTool).not.toHaveBeenCalled(); + }); + + it("scopes by serverId across plugins and retains describe ordering, deduplication, and metadata", async () => { + const f = fixture(); f.add("shared"); f.add("shared", "second-plugin"); f.add("unrelated"); + const gateway = f.createGateway(); + const tools = await gateway.listTools("shared"); + expect(tools).toHaveLength(2); + expect(new Set(tools.map((tool) => tool.opaqueId)).size).toBe(2); + const ids = [tools[1]!.opaqueId, "missing", tools[0]!.opaqueId, tools[1]!.opaqueId, "missing"]; + const result = describeTools(await gateway.resolveTools(ids), ids); + expect(result).toEqual({ tools: [tools[1], tools[0]], notFound: ["missing"] }); + expect(result.tools[0]).toMatchObject({ inputSchema: catalog().tools[0]!.inputSchema, outputSchema: catalog().tools[0]!.outputSchema, annotations: { readOnlyHint: true }, _meta: { fixture: true } }); + expect(searchTools(tools).servers).toHaveLength(2); + expect(f.starts()).toHaveLength(2); + }); + + it("coalesces concurrent owning-server discovery without retrying a failed mutation", async () => { + const f = fixture(); f.add("healthy"); + const gateway = f.createGateway(); + const [first, second] = await Promise.all([gateway.listTools("healthy"), gateway.listTools("healthy")]); + expect(second).toEqual(first); + expect(f.host.start).toHaveBeenCalledTimes(1); + f.host.callTool.mockRejectedValue(new Error("fixture mutation failed")); + await expect(gateway.call(first[0]!.opaqueId, {})).rejects.toThrow("fixture mutation failed"); + expect(f.host.callTool).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/bb-plugin-agent-plugins/src/gateway.ts b/packages/bb-plugin-agent-plugins/src/gateway.ts index 94d184e..c6cb50a 100644 --- a/packages/bb-plugin-agent-plugins/src/gateway.ts +++ b/packages/bb-plugin-agent-plugins/src/gateway.ts @@ -37,7 +37,8 @@ import type { import { optionalMcpCall } from "./mcp-compat.js"; export interface McpRuntime { - listTools(): Promise; + listTools(serverId?: string): Promise; + resolveTools(opaqueIds: string[]): Promise; call(opaqueId: string, args: JsonRecord, signal?: AbortSignal): Promise; listPrompts(): Promise; getPrompt(opaqueId: string, args?: JsonRecord, signal?: AbortSignal): Promise; @@ -115,6 +116,7 @@ const FETCH_TIMEOUT_GRACE_MS = 1_000; const CLOSE_TIMEOUT_MS = 2_000; const RETRY_AFTER_MS = 5_000; const CATALOG_TTL_MS = 5 * 60_000; +const CATALOG_CONCURRENCY = 4; function errorText(e: unknown): string { return e instanceof Error ? e.message : String(e); } function isOAuthStateMismatch(error: unknown): boolean { return errorText(error) === "MCP OAuth state mismatch"; } @@ -315,6 +317,8 @@ export class McpGateway implements McpRuntime { private readonly providers = new Map(); private readonly failures = new Map(); private readonly catalogCache = new Map(); + // Routing hints only: every use rechecks the server gates and current catalog. + private readonly toolOwners = new Map(); private readonly catalogLoads = new Map>(); private readonly catalogGenerations = new Map(); private readonly dirtyCatalogs = new Set(); @@ -353,6 +357,7 @@ export class McpGateway implements McpRuntime { this.conns.clear(); this.oauthPending.clear(); this.failures.clear(); + this.toolOwners.clear(); await Promise.all([...all.entries()].map(async ([key, value]) => { try { if ("expectedClose" in value) await this.closeConnected(key, value); @@ -711,32 +716,69 @@ export class McpGateway implements McpRuntime { finally { if (this.catalogLoads.get(key) === load) this.catalogLoads.delete(key); } } - async listTools(): Promise { - const result: CatalogTool[] = []; - for (const record of this.enabledApprovedServers()) { - const cfg = parseServerConfig(record.configJson); - if (!cfg) continue; - let catalog: CatalogCache; - try { catalog = await this.getCatalog(record); } - catch (error) { - this.log.warn(`listTools skip ${record.serverId}: ${errorText(error)}`); - result.push(this.toolError(record, cfg, error)); + async listTools(serverId?: string): Promise { + const scope = serverId?.trim(); + return this.toolsForServers(this.enabledApprovedServers().filter((record) => !scope || record.serverId === scope)); + } + + async resolveTools(opaqueIds: string[]): Promise { + const ids = new Set(opaqueIds); + const servers = this.enabledApprovedServers(); + const owners = new Set(); + for (const id of ids) { + const known = this.toolOwners.get(id); + if (known) { + owners.add(known); continue; } - const pluginName = this.store.getPlugin(record.pluginId)?.name ?? record.pluginId; - for (const tool of catalog.tools) result.push(this.catalogTool(record, cfg, pluginName, tool)); + // A cached ID can outlive this gateway instance. Use its stable prefix + // only to select a server, never as authority to execute a tool. Slugs + // are lossy (and may contain '__'), so ambiguous cold IDs fail closed; + // scoped discovery records the exact owner for those IDs. + if (!/^[a-z0-9_-]+_[a-f0-9]{10}$/u.test(id)) continue; + const candidates = servers.filter((record) => id.startsWith(`${slug(record.pluginId)}__${slug(record.serverId)}__`)); + if (candidates.length === 1) owners.add(keyOf(candidates[0]!.pluginId, candidates[0]!.serverId)); } - return result; + const catalog = await this.toolsForServers(servers.filter((record) => owners.has(keyOf(record.pluginId, record.serverId)))); + // Exact full-ID comparison (including the hash) rejects stale/forged IDs. + return catalog.filter((tool) => ids.has(tool.opaqueId)); + } + + private async toolsForServers(records: ReturnType): Promise { + const catalogs: CatalogTool[][] = new Array(records.length); + let next = 0; + // Per-request workers avoid a broad discovery queue blocking targeted calls. + // getCatalog still coalesces concurrent loads and backs off per server. + await Promise.all(Array.from({ length: Math.min(CATALOG_CONCURRENCY, records.length) }, async () => { + while (next < records.length) { + const index = next++; + const record = records[index]!; + const cfg = parseServerConfig(record.configJson); + if (!cfg) { catalogs[index] = []; continue; } + try { + const catalog = await this.getCatalog(record); + const pluginName = this.store.getPlugin(record.pluginId)?.name ?? record.pluginId; + catalogs[index] = catalog.tools.map((tool) => this.catalogTool(record, cfg, pluginName, tool)); + } catch (error) { + this.log.warn(`listTools skip ${record.serverId}: ${errorText(error)}`); + catalogs[index] = [this.toolError(record, cfg, error)]; + } + for (const tool of catalogs[index]!) this.toolOwners.set(tool.opaqueId, keyOf(record.pluginId, record.serverId)); + } + })); + // Preserve store/catalog order regardless of connection completion order. + return catalogs.flat(); } async call(opaqueId: string, args: JsonRecord, signal?: AbortSignal): Promise { - const catalog = await this.listTools(); + const catalog = await this.resolveTools([opaqueId]); const definition = catalog.find((item) => item.opaqueId === opaqueId); if (!definition) throw new Error(`Tool not found: ${opaqueId}`); if (definition.status === "error") throw new Error(definition.error ?? "tool in error state"); const conn = await this.ensureServer(definition.pluginId, definition.serverId); if (!conn) throw new Error(`MCP server unavailable: ${definition.serverId}`); const tool = conn.tools.find((item) => item.name === definition.name); + if (!tool) throw new Error(`Tool not found: ${opaqueId}`); const result = await this.withAuthState(definition.pluginId, definition.serverId, () => this.callTool(conn!, definition.name, args, tool, signal)); return contentResult(result); } @@ -924,6 +966,7 @@ export class McpGateway implements McpRuntime { const key = keyOf(pluginId, serverId); const record = this.serverRecord(pluginId, serverId); if (record.enabled !== 1) throw new Error(`MCP server ${serverId} is disabled`); + if (record.approved !== 1) throw new Error(`MCP server ${serverId} not approved (approval gate)`); const existing = this.conns.get(key); if (existing) return existing; if (this.oauthPending.has(key)) { @@ -932,7 +975,6 @@ export class McpGateway implements McpRuntime { } const inFlight = this.pending.get(key); if (inFlight) return inFlight; - if (record.approved !== 1) throw new Error(`MCP server ${serverId} not approved (approval gate)`); const failure = this.failures.get(key); if (failure && Date.now() - failure.at < RETRY_AFTER_MS) throw new Error(failure.message); const cfg = this.validatedConfig(record);