⚡ Bolt: Batch MCP OAuth token fetches to fix N+1 query

.jules/bolt.md

@@ -4,3 +4,6 @@
 ## 2026-05-16 - [Batch FeathersJS user fetches with $in operator to fix N+1 query]
 **Learning:** FeathersJS allows passing `$in` clauses through the query parameter (e.g. `user_id: { $in: ownerIds }`). When writing custom Feathers service logic, you can easily parse this array and pass it to Drizzle's `inArray()` to perform a batched query, instead of looping over `service.get(id)` causing N+1 database roundtrips.
 **Action:** When implementing or updating custom Feathers `find()` methods, extract and parse the `$in` parameters to support batched Drizzle `inArray()` lookups, and always replace `Promise.all(ids.map(id => service.get(id)))` with a single batched `find()` call.
+## 2025-07-20 - [Batch FeathersJS MCP OAuth token fetch to fix N+1 query]
+**Learning:** When a FeathersJS hook enriches multiple records by fetching related data (like per-user and shared OAuth tokens), performing individual database queries for each record inside a `Promise.all` loop causes an N+1 query performance bottleneck.
+**Action:** Pre-fetch all required related records using batched database queries with `inArray` based on the IDs of the records being enriched. Store the batched results in an in-memory Map using a compound key (e.g., `<server_id>:<user_id>` and `<server_id>:shared`), and look up the data from this Map during the loop.

apps/agor-daemon/src/register-hooks.ts

@@ -798,6 +798,47 @@ export function registerHooks(ctx: RegisterHooksContext): void {
       return context;
     }
 
+    let serversToProcess: MCPServer[] = [];
+    if (Array.isArray(context.result)) {
+      serversToProcess = context.result;
+    } else if (context.result?.data && Array.isArray(context.result.data)) {
+      serversToProcess = context.result.data;
+    } else if (context.result?.mcp_server_id) {
+      serversToProcess = [context.result];
+    }
+
+    const userTokenRepo = new UserMCPOAuthTokenRepository(db);
+
+    const perUserServers = serversToProcess.filter(
+      (s) => s.auth?.type === 'oauth' && (s.auth.oauth_mode ?? 'per_user') === 'per_user'
+    );
+    const sharedServers = serversToProcess.filter(
+      (s) => s.auth?.type === 'oauth' && s.auth.oauth_mode === 'shared'
+    );
+
+    const [perUserTokens, sharedTokens] = await Promise.all([
+      perUserServers.length > 0
+        ? userTokenRepo.getTokensForServers(
+            userId as import('@agor/core/types').UserID,
+            perUserServers.map((s) => s.mcp_server_id)
+          )
+        : Promise.resolve([]),
+      sharedServers.length > 0
+        ? userTokenRepo.getTokensForServers(
+            null,
+            sharedServers.map((s) => s.mcp_server_id)
+          )
+        : Promise.resolve([]),
+    ]);
+
+    const tokensMap = new Map<string, import('@agor/core/db').UserMCPOAuthToken>();
+    for (const token of perUserTokens) {
+      tokensMap.set(`${token.mcp_server_id}:${token.user_id}`, token);
+    }
+    for (const token of sharedTokens) {
+      tokensMap.set(`${token.mcp_server_id}:shared`, token);
+    }
+
     const injectToken = async (server: MCPServer) => {
       if (server.auth?.type !== 'oauth') {
         return server;
@@ -810,9 +851,13 @@ export function registerHooks(ctx: RegisterHooksContext): void {
       const tokenUserId: import('@agor/core/types').UserID | null =
         mode === 'per_user' ? (userId as import('@agor/core/types').UserID) : null;
 
+      const mapKey =
+        mode === 'per_user'
+          ? `${server.mcp_server_id}:${tokenUserId}`
+          : `${server.mcp_server_id}:shared`;
+
       try {
-        const userTokenRepo = new UserMCPOAuthTokenRepository(db);
-        const row = await userTokenRepo.getToken(tokenUserId, server.mcp_server_id);
+        const row = tokensMap.get(mapKey);
 
         if (!row) {
           console.log(

packages/core/src/db/repositories/user-mcp-oauth-tokens.ts

@@ -11,7 +11,7 @@
  */
 
 import type { MCPServerID, UserID } from '@agor/core/types';
-import { and, eq, isNull } from 'drizzle-orm';
+import { and, eq, inArray, isNull } from 'drizzle-orm';
 import type { Database } from '../client';
 import { deleteFrom, insert, select, update } from '../database-wrapper';
 import {
@@ -76,6 +76,35 @@ function matchKey(userId: UserID | null, serverId: MCPServerID) {
 export class UserMCPOAuthTokenRepository {
   constructor(private db: Database) {}
 
+  /**
+   * Batch look up token rows for multiple servers.
+   */
+  async getTokensForServers(
+    userId: UserID | null,
+    serverIds: MCPServerID[]
+  ): Promise<UserMCPOAuthToken[]> {
+    if (serverIds.length === 0) return [];
+    try {
+      const rows = await select(this.db)
+        .from(userMcpOauthTokens)
+        .where(
+          and(
+            inArray(userMcpOauthTokens.mcp_server_id, serverIds),
+            userId === null
+              ? isNull(userMcpOauthTokens.user_id)
+              : eq(userMcpOauthTokens.user_id, userId)
+          )
+        )
+        .all();
+      return rows.map(rowToToken);
+    } catch (error) {
+      throw new RepositoryError(
+        `Failed to get OAuth tokens: ${error instanceof Error ? error.message : String(error)}`,
+        error
+      );
+    }
+  }
+
   /**
    * Look up the token row for a (user, server) pair. Pass `null` for userId
    * to read the shared-mode row.