perf(sdk): batch invoke-path lookups with DataLoader semantics

.changeset/wild-buses-repeat.md

@@ -0,0 +1,11 @@
+---
+"executor": patch
+---
+
+Batch the SDK invoke path with DataLoader semantics. Concurrent `execute` calls
+in the same microtask window now share one storage query per table (tool,
+connection, integration) and one policy-rule-set snapshot instead of four point
+queries per call, so naive fan-out code (`Promise.all` over tool calls, code
+mode loops, parallel MCP tool calls) no longer produces N+1 query storms.
+Sequential calls are unchanged, and transactional reads bypass the batch window
+so transaction isolation is preserved.

packages/core/sdk/src/executor.ts

@@ -147,6 +147,7 @@ import {
   type OAuthEndpointUrlPolicy,
 } from "./oauth-helpers";
 import { connectionIdentifier } from "./connection-name-identifier";
+import { makeInvokeBatching } from "./invoke-batching";
 import { annotateToolResultOutcome } from "./tool-result";
 
 const PLUGIN_STORAGE_DELETE_KEY_BATCH_SIZE = 90;
@@ -2532,6 +2533,51 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea
               ),
             );
 
+    // ------------------------------------------------------------------
+    // Invoke-path batching. Concurrent `execute` calls within one microtask
+    // window share ONE query per table (tool / connection / integration) and
+    // one policy-rule-set snapshot instead of 4×N point queries — the
+    // DataLoader pattern, via Effect Request/RequestResolver. Loaders serve
+    // a batch of mixed tuples with a single OR-of-tuples query; the memory
+    // and SQL adapters both plan it as one round trip.
+    // ------------------------------------------------------------------
+
+    const invokeBatching = makeInvokeBatching({
+      loadToolRows: (requests) =>
+        core.findMany("tool", {
+          where: (b: AnyCb) =>
+            b.or(
+              ...requests.map((request) =>
+                b.and(
+                  byOwner(request.owner)(b),
+                  b("integration", "=", String(request.integration)),
+                  b("connection", "=", String(request.connection)),
+                  b("name", "=", String(request.tool)),
+                ),
+              ),
+            ),
+          select: TOOL_INVOCATION_COLUMNS,
+        }),
+      loadConnectionRows: (requests) =>
+        core.findMany("connection", {
+          where: (b: AnyCb) =>
+            b.or(
+              ...requests.map((request) =>
+                b.and(
+                  byOwner(request.owner)(b),
+                  b("integration", "=", String(request.integration)),
+                  b("name", "=", String(request.name)),
+                ),
+              ),
+            ),
+        }),
+      loadIntegrationRows: (slugs) =>
+        core.findMany("integration", {
+          where: (b: AnyCb) => b("slug", "in", slugs.map(String)),
+        }),
+      loadPolicyRuleSet: () => listActivePolicyRuleSet(),
+    });
+
     // ------------------------------------------------------------------
     // Tools (read surface)
     // ------------------------------------------------------------------
@@ -3035,7 +3081,7 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea
         // not the 5-segment dynamic form.
         const staticEntry = staticTools.get(String(address));
         if (staticEntry) {
-          const policyRules = yield* listActivePolicyRuleSet();
+          const policyRules = yield* invokeBatching.getPolicyRuleSet();
           const policy = yield* resolvePolicyFromRuleSet(
             String(address),
             policyRules,
@@ -3064,16 +3110,13 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea
 
         // Find the tool row — projected: invoke needs routing/policy fields
         // only, never the multi-KB input/output schema JSON (`tools.schema`
-        // is the schema-bearing surface).
-        const row = yield* core.findFirst("tool", {
-          where: (b: AnyCb) =>
-            b.and(
-              byOwner(parsed.owner)(b),
-              b("integration", "=", String(parsed.integration)),
-              b("connection", "=", String(parsed.connection)),
-              b("name", "=", String(parsed.tool)),
-            ),
-          select: TOOL_INVOCATION_COLUMNS,
+        // is the schema-bearing surface). Batched: concurrent executes in the
+        // same window share one query.
+        const row = yield* invokeBatching.getToolRow({
+          owner: parsed.owner,
+          integration: parsed.integration,
+          connection: parsed.connection,
+          tool: parsed.tool,
         });
         if (!row) {
           const searchMatches = yield* searchToolRowsForConnection(parsed);
@@ -3087,7 +3130,7 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea
 
         // Resolve policy (owner-ranked).
         const toolForPolicy = rowToTool(row);
-        const policyRules = yield* listActivePolicyRuleSet();
+        const policyRules = yield* invokeBatching.getPolicyRuleSet();
         const annotations = decodeJsonColumn(row.annotations) as ToolAnnotations | undefined;
         const policy = yield* resolvePolicyFromRuleSet(
           normalizedPolicyId(toolForPolicy),
@@ -3115,8 +3158,8 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea
           });
         }
 
-        // Find the connection row.
-        const connectionRow = yield* findConnectionRow({
+        // Find the connection row (batched with peer executes).
+        const connectionRow = yield* invokeBatching.getConnectionRow({
           owner: parsed.owner,
           integration: parsed.integration,
           name: parsed.connection,
@@ -3147,7 +3190,7 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea
         // Resolve every named credential input (`variable → value`); `value` is
         // the primary `token` for single-input + OAuth callers.
         const values = yield* resolveConnectionValues(connectionRow);
-        const integrationRow = yield* findIntegrationRow(parsed.integration);
+        const integrationRow = yield* invokeBatching.getIntegrationRow(parsed.integration);
         const credential: ToolInvocationCredential = {
           owner: parsed.owner,
           integration: parsed.integration,

packages/core/sdk/src/invoke-batching.test.ts

@@ -0,0 +1,201 @@
+import { describe, expect, it } from "@effect/vitest";
+import { Effect, Predicate, Result } from "effect";
+
+import { createExecutor, collectTables } from "./executor";
+import {
+  AuthTemplateSlug,
+  ConnectionName,
+  IntegrationSlug,
+  ProviderItemId,
+  ProviderKey,
+  Subject,
+  Tenant,
+  ToolAddress,
+  ToolName,
+} from "./ids";
+import { definePlugin } from "./plugin";
+import type { CredentialProvider } from "./provider";
+import { createSqliteTestFumaDb } from "./sqlite-test-db";
+import { withQueryContext } from "@executor-js/fumadb/query";
+import type { FumaDb } from "./fuma-runtime";
+
+// ---------------------------------------------------------------------------
+// Query-counting FumaDb wrapper. Counts reads per `${method}:${table}` while
+// forwarding everything else, and re-wraps `withContext` results so the
+// counter survives the executor's own context application.
+// ---------------------------------------------------------------------------
+
+type QueryCounts = Map<string, number>;
+
+const COUNTED_METHODS = new Set(["findFirst", "findMany", "count"]);
+
+const countingDb = (db: FumaDb, counts: QueryCounts): FumaDb => {
+  // oxlint-disable-next-line executor/no-double-cast -- boundary: test spy walks the ORM object's own members reflectively
+  const record = db as unknown as Record<string, unknown>;
+  const wrapper: Record<string, unknown> = {};
+  // Copy every member (methods included) off the real query object; a Proxy
+  // can't stand in because FumaDb's properties are non-configurable (proxy
+  // invariant violation on `get`). `withContext` and friends are
+  // non-enumerable, so walk own property names, not Object.keys.
+  for (const key of Object.getOwnPropertyNames(record)) {
+    const value = record[key];
+    if (typeof value !== "function") {
+      wrapper[key] = value;
+      continue;
+    }
+    const fn = value as (...a: unknown[]) => unknown;
+    if (COUNTED_METHODS.has(key)) {
+      wrapper[key] = (table: string, ...rest: unknown[]) => {
+        const countKey = `${key}:${table}`;
+        counts.set(countKey, (counts.get(countKey) ?? 0) + 1);
+        return fn.call(db, table, ...rest);
+      };
+    } else if (key === "withContext") {
+      wrapper[key] = (context: unknown) =>
+        countingDb((fn as (c: unknown) => FumaDb).call(db, context), counts);
+    } else {
+      wrapper[key] = (...args: unknown[]) => fn.call(db, ...args);
+    }
+  }
+  // oxlint-disable-next-line executor/no-double-cast -- boundary: the wrapper delegates every FumaDb member 1:1
+  return wrapper as unknown as FumaDb;
+};
+
+const totalFor = (counts: QueryCounts, table: string): number => {
+  let total = 0;
+  for (const [key, count] of counts) {
+    if (key.endsWith(`:${table}`)) total += count;
+  }
+  return total;
+};
+
+// ---------------------------------------------------------------------------
+// Fixture: a plugin with one integration and two tools.
+// ---------------------------------------------------------------------------
+
+const INTEG = IntegrationSlug.make("demo");
+const CONN = ConnectionName.make("main");
+
+const memoryProvider = (): CredentialProvider => {
+  const store = new Map<string, string>();
+  return {
+    key: ProviderKey.make("memory"),
+    writable: true,
+    get: (id) => Effect.sync(() => store.get(String(id)) ?? null),
+    set: (id, value) => Effect.sync(() => void store.set(String(id), value)),
+  };
+};
+
+const demoPlugin = definePlugin(() => ({
+  id: "demo" as const,
+  credentialProviders: [memoryProvider()],
+  storage: () => ({}),
+  resolveTools: () =>
+    Effect.succeed({
+      tools: [
+        { name: ToolName.make("alpha"), description: "alpha" },
+        { name: ToolName.make("beta"), description: "beta" },
+      ],
+    }),
+  invokeTool: ({ toolRow, args }) => Effect.succeed({ ran: toolRow.name, args }),
+  extension: (ctx) => ({
+    seed: () => ctx.core.integrations.register({ slug: INTEG, description: "Demo", config: {} }),
+  }),
+}))();
+
+const addr = (tool: string): ToolAddress => ToolAddress.make(`tools.${INTEG}.org.${CONN}.${tool}`);
+
+const TENANT = "test-tenant";
+const SUBJECT = "test-subject";
+
+const makeCountedExecutor = Effect.fnUntraced(function* () {
+  const counts: QueryCounts = new Map();
+  const tables = collectTables();
+  const testDb = yield* Effect.promise(() =>
+    createSqliteTestFumaDb({ tables, namespace: "executor_test" }),
+  );
+  const db = withQueryContext(countingDb(testDb.db, counts), {
+    tenant: TENANT,
+    subject: SUBJECT,
+  });
+  const executor = yield* createExecutor({
+    tenant: Tenant.make(TENANT),
+    subject: Subject.make(SUBJECT),
+    db,
+    plugins: [demoPlugin] as const,
+    onElicitation: "accept-all",
+  });
+  yield* executor.demo.seed();
+  yield* executor.connections.create({
+    owner: "org",
+    name: CONN,
+    integration: INTEG,
+    template: AuthTemplateSlug.make("apiKey"),
+    from: { provider: ProviderKey.make("memory"), id: ProviderItemId.make("v") },
+  });
+  return { executor, counts };
+});
+
+describe("invoke-path batching", () => {
+  it.effect("N concurrent executes share one query per table (no N+1)", () =>
+    Effect.gen(function* () {
+      const { executor, counts } = yield* makeCountedExecutor();
+
+      counts.clear();
+      const N = 25;
+      const results = yield* Effect.all(
+        Array.from({ length: N }, (_, i) =>
+          executor.execute(addr(i % 2 === 0 ? "alpha" : "beta"), { i }),
+        ),
+        { concurrency: "unbounded" },
+      );
+
+      expect(results).toHaveLength(N);
+      expect(results[0]).toEqual({ ran: "alpha", args: { i: 0 } });
+
+      // The batch window collapses all N lookups into one query per table.
+      // Allow a little slack (the runtime may split across two microtask
+      // windows under load) but the point is O(1), not O(N).
+      expect(totalFor(counts, "tool")).toBeLessThanOrEqual(2);
+      expect(totalFor(counts, "connection")).toBeLessThanOrEqual(2);
+      expect(totalFor(counts, "integration")).toBeLessThanOrEqual(2);
+      expect(totalFor(counts, "tool_policy")).toBeLessThanOrEqual(2);
+    }),
+  );
+
+  it.effect("sequential executes still behave (batch of one)", () =>
+    Effect.gen(function* () {
+      const { executor, counts } = yield* makeCountedExecutor();
+
+      counts.clear();
+      const first = yield* executor.execute(addr("alpha"), { seq: 1 });
+      const second = yield* executor.execute(addr("beta"), { seq: 2 });
+
+      expect(first).toEqual({ ran: "alpha", args: { seq: 1 } });
+      expect(second).toEqual({ ran: "beta", args: { seq: 2 } });
+      // Two sequential invokes = two windows: exactly the per-call point
+      // queries the unbatched path made, no regression.
+      expect(totalFor(counts, "tool")).toBe(2);
+      expect(totalFor(counts, "connection")).toBe(2);
+    }),
+  );
+
+  it.effect("a missing tool inside a batch fails alone, peers succeed", () =>
+    Effect.gen(function* () {
+      const { executor } = yield* makeCountedExecutor();
+
+      const [ok, missing] = yield* Effect.all(
+        [
+          Effect.result(executor.execute(addr("alpha"), {})),
+          Effect.result(executor.execute(addr("nope"), {})),
+        ],
+        { concurrency: "unbounded" },
+      );
+
+      expect(Result.isSuccess(ok)).toBe(true);
+      expect(Result.isFailure(missing)).toBe(true);
+      if (!Result.isFailure(missing)) return;
+      expect(Predicate.isTagged(missing.failure, "ToolNotFoundError")).toBe(true);
+    }),
+  );
+});