Narrow a Bot's tools in the query, not after reading every one (#184)

`listForAgent` selected every row in `mcp_tools` and then discarded the ones the
Bot was not granted, with a linear `includes` over the grants per row. It is the
run-time path: it runs on every run of every Bot, before anything else, and it
sits underneath the tool selection added in #178 — so its cost is paid first, on
exactly the large catalogues that change exists to make work. At the thousand
tools #119 names as the target, that is a thousand rows across the wire and a
thousand walks of the grant list to offer a handful.

The query now names the servers the Bot holds something from, which is a
predicate the `(server_id, name)` primary key can use, and the grant list is a
Set rather than an array. `knownToolRefs` in this same file already did it this
way, with the same reasoning written above it; this is the run-time half catching
up with the save-time half.

Narrowing by server rather than by exact pair on purpose. A clause per grant
would be exact, and a server's own tool list is already the bound on what comes
back, so it buys little for a where clause that grows with the grants. The exact
ref is still matched after the read, and that match is now load-bearing in a way
it was not before: narrowing by server alone would offer every tool of any
server the Bot holds anything from. There is a test for that, because every
other test in the file passes without it.

Behaviour is unchanged. Same tools, same order, same refusals.

Author: Hotragn

server/src/plugins/store.ts

@@ -1050,12 +1050,34 @@ export function createPluginStore(options: PluginStoreOptions) {
         .filter((row) => row.kind === "skill")
         .map((row) => row.ref);
 
+      /*
+       * Narrowed in the query to the servers this Bot is actually granted something from, the same way
+       * `knownToolRefs` does it and for the same reason: a deployment aiming at a thousand tools should
+       * not read all of them to offer a handful. This is the run-time path, so it ran on every run of
+       * every Bot, selected every row in `mcp_tools`, and then discarded almost all of them here — and
+       * it sits underneath tool selection, so its cost is paid before the narrowing that was added to
+       * make large catalogues work.
+       *
+       * The exact ref is still matched below rather than in the query. Narrowing by server is a
+       * predicate the composite primary key can use; naming every (server, tool) pair would be exact
+       * and is not worth a clause per grant, because a server's own tool list is the bound on what
+       * comes back.
+       */
+      const grantedServers = [
+        ...new Set(toolRefs.map((ref) => ref.split("/")[0] ?? "")),
+      ];
       const toolRows =
-        toolRefs.length === 0
+        grantedServers.length === 0
           ? []
-          : await database.select().from(mcpTools).orderBy(asc(mcpTools.name));
+          : await database
+              .select()
+              .from(mcpTools)
+              .where(inArray(mcpTools.serverId, grantedServers))
+              .orderBy(asc(mcpTools.name));
+      // A set, so this is a lookup per row rather than a walk of the grants per row.
+      const granted = new Set(toolRefs);
       const grantedTools = toolRows
-        .filter((row) => toolRefs.includes(`${row.serverId}/${row.name}`))
+        .filter((row) => granted.has(`${row.serverId}/${row.name}`))
         .map((row) => {
           const ref = `${row.serverId}/${row.name}`;
           return {

server/tests/plugin-store.integration.test.ts

@@ -35,6 +35,8 @@ const strangerId = `agent_plugin_stranger_${suite}`;
 const serverId = "google-drive";
 const toolName = "search_files";
 const ref = `${serverId}/${toolName}`;
+/** A tool on the same server that nobody is granted. Suite-scoped, so it is never a real one. */
+const siblingToolName = `not_granted_${suite}`;
 
 let policy: ActionPolicy = { mode: "enforce", deny: [], allow: ["true"] };
 
@@ -136,6 +138,22 @@ beforeAll(async () => {
     .insert(mcpTools)
     .values({ serverId, name: toolName, description: "Search files." })
     .onConflictDoNothing();
+  /*
+   * A second tool on the SAME server, granted to nobody.
+   *
+   * `listForAgent` narrows to the servers a Bot holds something from and then matches the exact ref,
+   * and this is what makes the second half load-bearing: without it, holding one tool from a server
+   * would offer every tool that server has. Suite-scoped, so it is unambiguously a fixture and
+   * cannot collide with a name the vendor really advertises.
+   */
+  await database
+    .insert(mcpTools)
+    .values({
+      serverId,
+      name: siblingToolName,
+      description: "A tool on the same server that nobody was granted.",
+    })
+    .onConflictDoNothing();
 });
 
 afterAll(async () => {
@@ -157,6 +175,12 @@ afterAll(async () => {
         inArray(pluginGrants.agentId, [holderId, strangerId]),
       ),
     );
+  // Suite-scoped, so it is this suite's whatever else is true of the server.
+  await database
+    .delete(mcpTools)
+    .where(
+      and(eq(mcpTools.serverId, serverId), eq(mcpTools.name, siblingToolName)),
+    );
   // A server row is deployment configuration, so it belongs to the deployment rather than here.
   // The fixture tool goes whether or not this suite owns the server, but only if it put it there.
   if (!toolWasAlreadyAdvertised) {
@@ -219,6 +243,22 @@ describe("a grant is the permission", () => {
     expect(nothing.tools).toEqual([]);
     expect(nothing.skills).toEqual([]);
   });
+
+  test("holding one tool from a server does not offer that server's others", async () => {
+    /*
+     * The property the exact-ref match protects, now that the query narrows by server rather than
+     * reading the whole catalogue. Widening this to "every tool on a server you hold anything from"
+     * would pass every other test in this file: the Bot would still be offered what it holds, and the
+     * stranger would still be offered nothing.
+     */
+    await store.grant("mcp", ref, holderId, "admin@openbot.local");
+    const held = await store.listForAgent(holderId);
+
+    expect(held.tools.map((tool) => tool.ref)).toEqual([ref]);
+    expect(held.tools.map((tool) => tool.ref)).not.toContain(
+      `${serverId}/${siblingToolName}`,
+    );
+  });
 });
 
 describe("the policy is asked as well as the grant", () => {