tools: collapse annotation N+1 + push sourceId filter to DB

openapi/graphql resolveAnnotations ran one listOperationsBySource per distinct scope in the toolRows. Store the owning scopeId on each StoredOperation instead and issue a single query, partitioning by op.scopeId so shadowed tools still resolve against the right binding. Also push executor.tools.list's sourceId filter down to storage so `/sources/:id/tools` stops pulling every row in the scope stack.

Rhys Sullivan committed 6 months ago

Commit 400a8061 parent aa220d75 files · +44 −40Browse files

Showing the changes to packages/plugins/openapi/src (2 of 5 files). Show all

Unified · Split

2 files

packages/plugins/openapi/src/sdk/plugin.ts +13 −12
packages/plugins/openapi/src/sdk/store.ts +4 −6

packages/plugins/openapi/src/sdk/plugin.ts+13−12

         const storedOps: StoredOperation[] = definitions.map((def) => ({
           toolId: `${namespace}.${def.toolPath}`,
           sourceId: namespace,
+          scopeId: input.scope,
           binding: toBinding(def),
         }));
           // toolRows for a single (plugin_id, source_id) group can still
           // straddle multiple scopes when the source is shadowed (e.g. an
           // org-level openapi source plus a per-user override that
-          // re-registers the same tool ids). Run one listOperationsBySource
+          // re-registers the same tool ids). One query returns every op
-          // per distinct scope so each lookup pins {source_id, scope_id}
+          // row in the caller's stack (the scoped adapter pins
-          // and we don't fall through to the wrong scope's bindings.
+          // `scope_id IN (stack)`); partition by `op.scopeId` so each
-          const scopes = new Set<string>();
+          // tool row resolves against its own scope's binding.
-          for (const row of toolRows as readonly ToolRow[]) {
+          const ops = yield* ctx.storage.listOperationsBySource(sourceId);
-            scopes.add(row.scope_id as string);
-          }
           const byScope = new Map<string, Map<string, OperationBinding>>();
-          for (const scope of scopes) {
+          for (const op of ops) {
-            const ops = yield* ctx.storage.listOperationsBySource(sourceId, scope);
+            let byId = byScope.get(op.scopeId);
-            const byId = new Map<string, OperationBinding>();
+            if (!byId) {
-            for (const op of ops) byId.set(op.toolId, op.binding);
+              byId = new Map<string, OperationBinding>();
-            byScope.set(scope, byId);
+              byScope.set(op.scopeId, byId);
+            }
+            byId.set(op.toolId, op.binding);
           }
           const out: Record<string, ToolAnnotations> = {};

packages/plugins/openapi/src/sdk/store.ts+4−6

 export interface StoredOperation {
   readonly toolId: string;
   readonly sourceId: string;
+  readonly scopeId: string;
   readonly binding: OperationBinding;
 }
   readonly listOperationsBySource: (
     sourceId: string,
-    scope: string,
   ) => Effect.Effect<readonly StoredOperation[], StorageFailure>;
   readonly removeSource: (
   const rowToOperation = (row: Record<string, unknown>): StoredOperation => ({
     toolId: row.id as string,
     sourceId: row.source_id as string,
+    scopeId: row.scope_id as string,
     binding: decodeBinding(
       typeof row.binding === "string" ? JSON.parse(row.binding) : row.binding,
     ),
         })
         .pipe(Effect.map((row) => (row ? rowToOperation(row) : null))),
-    listOperationsBySource: (sourceId, scope) =>
+    listOperationsBySource: (sourceId) =>
       adapter
         .findMany({
           model: "openapi_operation",
-          where: [
+          where: [{ field: "source_id", value: sourceId }],
-            { field: "source_id", value: sourceId },
-            { field: "scope_id", value: scope },
-          ],
         })
         .pipe(Effect.map((rows) => rows.map(rowToOperation))),