From 5dd886c3494c289f4bcedee3d4657a68df9d4b1d Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:31:54 -0700 Subject: [PATCH] perf(sdk): batch invoke-path lookups with Effect RequestResolver Concurrent execute calls in the same microtask window now share one query per table (tool, connection, integration) plus one policy rule-set snapshot instead of four point queries per call, killing the N+1 pattern under code-mode fan-out and parallel MCP tool calls. - invoke-batching.ts: Request types + resolvers keyed by plain data; batch loaders serve mixed tuples with a single OR-of-tuples query and dedupe equal keys before hitting storage. No cross-batch caching, so read-your-writes behavior is unchanged. - Transactional callers bypass the batch window (a shared batch fiber would read through the wrong activeFumaDbRef) and keep exact per-transaction semantics. - Proof test counts real queries through a spying FumaDb wrapper: 25 concurrent executes touch each table at most twice; sequential calls still issue the same per-call point reads as before. --- .changeset/wild-buses-repeat.md | 11 + packages/core/sdk/src/executor.ts | 73 ++++- packages/core/sdk/src/invoke-batching.test.ts | 201 ++++++++++++ packages/core/sdk/src/invoke-batching.ts | 303 ++++++++++++++++++ 4 files changed, 573 insertions(+), 15 deletions(-) create mode 100644 .changeset/wild-buses-repeat.md create mode 100644 packages/core/sdk/src/invoke-batching.test.ts create mode 100644 packages/core/sdk/src/invoke-batching.ts diff --git a/.changeset/wild-buses-repeat.md b/.changeset/wild-buses-repeat.md new file mode 100644 index 000000000..67aa93af8 --- /dev/null +++ b/.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. diff --git a/packages/core/sdk/src/executor.ts b/packages/core/sdk/src/executor.ts index 7f46b2279..ae3d914ac 100644 --- a/packages/core/sdk/src/executor.ts +++ b/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 = + 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 = - 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 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; + const wrapper: Record = {}; + // 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(); + 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); + }), + ); +}); diff --git a/packages/core/sdk/src/invoke-batching.ts b/packages/core/sdk/src/invoke-batching.ts new file mode 100644 index 000000000..8f84ce46a --- /dev/null +++ b/packages/core/sdk/src/invoke-batching.ts @@ -0,0 +1,303 @@ +// --------------------------------------------------------------------------- +// Invoke-path batching: DataLoader semantics via Effect Request/RequestResolver. +// +// The invoke hot path loads four things per call: the tool row, the active +// policy rule set, the connection row, and the integration row. Called once, +// that's four point queries, which is fine. Called N times concurrently (code +// mode fanning out `await Promise.all(items.map(x => tools.foo.bar(x)))`, an +// MCP host serving parallel tool calls), it's 4×N queries against the same +// handful of rows: the classic N+1. +// +// Each lookup below is an Effect `Request` plus a `RequestResolver`. The +// Effect runtime collects every request issued within the same microtask +// window (resolver delay is `Effect.yieldNow`) and hands the resolver the +// whole batch, which it serves with ONE query per table. Sequential callers +// pay nothing: a batch of one is exactly the point query this replaced. There +// is no cross-batch caching; every batch re-reads storage, so invalidation +// semantics are unchanged. +// +// Requests are keyed by plain data (owner/slug/name strings), never rows, so +// equality is meaningful and nothing heavy is retained. +// --------------------------------------------------------------------------- + +import { Effect, Exit, Request, RequestResolver } from "effect"; + +import type { ConnectionRow, IntegrationRow, ToolInvocationRow } from "./core-schema"; +import { activeFumaDbRef, type StorageFailure } from "./fuma-runtime"; +import type { ConnectionName, IntegrationSlug, Owner, ToolName } from "./ids"; + +// --------------------------------------------------------------------------- +// Request types +// --------------------------------------------------------------------------- + +interface GetToolRow extends Request.Request { + readonly _tag: "GetToolRow"; + readonly owner: Owner; + readonly integration: IntegrationSlug; + readonly connection: ConnectionName; + readonly tool: ToolName; +} +const GetToolRow = Request.tagged("GetToolRow"); + +interface GetConnectionRow extends Request.Request { + readonly _tag: "GetConnectionRow"; + readonly owner: Owner; + readonly integration: IntegrationSlug; + readonly name: ConnectionName; +} +const GetConnectionRow = Request.tagged("GetConnectionRow"); + +interface GetIntegrationRow extends Request.Request { + readonly _tag: "GetIntegrationRow"; + readonly slug: IntegrationSlug; +} +const GetIntegrationRow = Request.tagged("GetIntegrationRow"); + +// The policy rule set is one shared snapshot per batch window: every request +// is identical, so the resolver runs the underlying load once and fans the +// result out to all awaiting fibers. +interface GetPolicyRuleSet extends Request.Request { + readonly _tag: "GetPolicyRuleSet"; +} + +// --------------------------------------------------------------------------- +// Wiring: the executor passes its storage closures in; we hand back +// point-lookup functions with identical signatures to the ones they replace. +// --------------------------------------------------------------------------- + +export interface InvokeBatchingDeps { + /** Batched tool-row load: one query for all (owner, integration, connection, + * tool) tuples in the window, projected to the invocation columns. */ + readonly loadToolRows: ( + requests: readonly { + readonly owner: Owner; + readonly integration: IntegrationSlug; + readonly connection: ConnectionName; + readonly tool: ToolName; + }[], + ) => Effect.Effect; + readonly loadConnectionRows: ( + requests: readonly { + readonly owner: Owner; + readonly integration: IntegrationSlug; + readonly name: ConnectionName; + }[], + ) => Effect.Effect; + readonly loadIntegrationRows: ( + slugs: readonly IntegrationSlug[], + ) => Effect.Effect; + readonly loadPolicyRuleSet: () => Effect.Effect; +} + +export interface InvokeBatching { + readonly getToolRow: (input: { + readonly owner: Owner; + readonly integration: IntegrationSlug; + readonly connection: ConnectionName; + readonly tool: ToolName; + }) => Effect.Effect; + readonly getConnectionRow: (input: { + readonly owner: Owner; + readonly integration: IntegrationSlug; + readonly name: ConnectionName; + }) => Effect.Effect; + readonly getIntegrationRow: ( + slug: IntegrationSlug, + ) => Effect.Effect; + readonly getPolicyRuleSet: () => Effect.Effect; +} + +// Batches from different fibers execute the `runAll` effect on a fiber forked +// with the FIRST caller's context. Inside `fuma.transaction` the storage +// handle rides on `activeFumaDbRef` in fiber context, so joining a shared +// batch would read through the wrong (or no) transaction. Transactional +// callers bypass the batch window and run the load directly on their own +// fiber — same query, exact transaction semantics, no cross-fiber sharing. +const inTransaction: Effect.Effect = Effect.map( + Effect.service(activeFumaDbRef), + (active) => active !== null, +); + +/** Complete every entry in a batch from a keyed index of loaded rows. A + * request whose key has no row completes with `null`: absence is a value + * (the invoke path turns it into its own typed not-found error). A load + * failure fails every entry in the batch, mirroring what each point query + * would have done. */ +const completeFromIndex = >( + entries: readonly Request.Entry[], + load: Effect.Effect>, StorageFailure>, + keyOf: (request: A) => string, +): Effect.Effect => + load.pipe( + Effect.matchEffect({ + onSuccess: (index) => + Effect.sync(() => { + for (const entry of entries) { + entry.completeUnsafe( + Exit.succeed((index.get(keyOf(entry.request)) ?? null) as Request.Success), + ); + } + }), + onFailure: (error) => + Effect.sync(() => { + for (const entry of entries) { + // The constraint pins A's error channel to StorageFailure, but the + // conditional `Request.Error` doesn't reduce over an unresolved + // generic — assert the equivalence once here. + entry.completeUnsafe(Exit.fail(error) as Exit.Exit>); + } + }), + }), + ); + +/** The runtime collects entries per window without deduping equal requests; + * two fibers asking for the same connection contribute two entries. Loaders + * receive each distinct key once (DataLoader's contract) and both entries + * complete from the shared index. */ +const uniqueBy = (items: readonly T[], key: (item: T) => string): readonly T[] => { + const seen = new Set(); + const out: T[] = []; + for (const item of items) { + const k = key(item); + if (seen.has(k)) continue; + seen.add(k); + out.push(item); + } + return out; +}; + +const toolKey = (r: { + readonly owner: string; + readonly integration: string; + readonly connection: string; + readonly tool: string; +}): string => `${r.owner} ${r.integration} ${r.connection} ${r.tool}`; + +const connectionKey = (r: { + readonly owner: string; + readonly integration: string; + readonly name: string; +}): string => `${r.owner} ${r.integration} ${r.name}`; + +export const makeInvokeBatching = ( + deps: InvokeBatchingDeps, +): InvokeBatching => { + const toolResolver = RequestResolver.make((entries) => + completeFromIndex( + entries, + deps + .loadToolRows( + uniqueBy( + entries.map((entry) => entry.request), + toolKey, + ), + ) + .pipe( + Effect.map( + (rows) => + new Map( + rows.map((row) => [ + toolKey({ + owner: row.owner, + integration: row.integration, + connection: row.connection, + tool: row.name, + }), + row, + ]), + ), + ), + ), + toolKey, + ), + ); + + const connectionResolver = RequestResolver.make((entries) => + completeFromIndex( + entries, + deps + .loadConnectionRows( + uniqueBy( + entries.map((entry) => entry.request), + connectionKey, + ), + ) + .pipe( + Effect.map( + (rows) => + new Map( + rows.map((row) => [ + connectionKey({ + owner: row.owner, + integration: row.integration, + name: row.name, + }), + row, + ]), + ), + ), + ), + connectionKey, + ), + ); + + const integrationResolver = RequestResolver.make((entries) => + completeFromIndex( + entries, + deps + .loadIntegrationRows( + uniqueBy( + entries.map((entry) => entry.request.slug), + String, + ), + ) + .pipe(Effect.map((rows) => new Map(rows.map((row) => [String(row.slug), row])))), + (request) => String(request.slug), + ), + ); + + const policyResolver = RequestResolver.make>((entries) => + deps.loadPolicyRuleSet().pipe( + Effect.matchEffect({ + onSuccess: (ruleSet) => + Effect.sync(() => { + for (const entry of entries) entry.completeUnsafe(Exit.succeed(ruleSet)); + }), + onFailure: (error) => + Effect.sync(() => { + for (const entry of entries) entry.completeUnsafe(Exit.fail(error)); + }), + }), + ), + ); + const policyRequest = Request.of>()({ + _tag: "GetPolicyRuleSet", + }); + + const first = (rows: readonly A[]): A | null => rows[0] ?? null; + + return { + getToolRow: (input) => + Effect.flatMap(inTransaction, (transactional) => + transactional + ? Effect.map(deps.loadToolRows([input]), first) + : Effect.request(GetToolRow(input), toolResolver), + ), + getConnectionRow: (input) => + Effect.flatMap(inTransaction, (transactional) => + transactional + ? Effect.map(deps.loadConnectionRows([input]), first) + : Effect.request(GetConnectionRow(input), connectionResolver), + ), + getIntegrationRow: (slug) => + Effect.flatMap(inTransaction, (transactional) => + transactional + ? Effect.map(deps.loadIntegrationRows([slug]), first) + : Effect.request(GetIntegrationRow({ slug }), integrationResolver), + ), + getPolicyRuleSet: () => + Effect.flatMap(inTransaction, (transactional) => + transactional ? deps.loadPolicyRuleSet() : Effect.request(policyRequest, policyResolver), + ), + }; +};