⚡ Bolt: Batch MCP OAuth Token resolution

.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.
+## 2024-07-30 - [Batch Fetch MCP OAuth Tokens to eliminate N+1 queries]
+**Learning:** When fetching records across multiple sources and permissions, like MCP OAuth tokens for shared and per-user configurations simultaneously in FeathersJS hooks, making individual DB calls within a `Promise.all` causes N+1 problems. Instead, use a compound key for an in-memory Map (e.g. `<server_id>:<user_id>`) from a single batched database query resolving both types.
+**Action:** When populating an in-memory Map from a batched database result to eliminate N+1 queries, use a compound key combining item ID and access level identifier (like user ID) to preserve both shared and exclusive records, avoiding premature squashing or missing records that would otherwise re-trigger N+1 queries on Map misses.

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

@@ -1,3 +1,4 @@
+import type { UserMCPOAuthToken } from '@agor/core/db';
 /**
  * Service Hooks Registration
  *
@@ -788,16 +789,45 @@ export function registerHooks(ctx: RegisterHooksContext): void {
       : queryForUserId
         ? 'query-param'
         : 'none';
+
+    let servers: MCPServer[] = [];
+    if (Array.isArray(context.result)) {
+      servers = context.result;
+    } else if (context.result?.data && Array.isArray(context.result.data)) {
+      servers = context.result.data;
+    } else if (context.result?.mcp_server_id) {
+      servers = [context.result];
+    }
+
     console.log(
       `[MCP OAuth] injectPerUserOAuthTokens called - userId: ${userId || 'NONE'}, ` +
         `source: ${source}, provider: ${context.params?.provider || 'internal'}, ` +
-        `method: ${context.method}, resultCount: ${Array.isArray(context.result) ? context.result.length : 1}`
+        `method: ${context.method}, resultCount: ${servers.length}`
     );
     if (!userId) {
       console.log('[MCP OAuth] No user ID - skipping token injection');
       return context;
     }
 
+    const userTokenRepo = new UserMCPOAuthTokenRepository(db);
+    const tokenMap = new Map<string, UserMCPOAuthToken>();
+
+    const serverIds = servers
+      .filter((s) => s.auth?.type === 'oauth')
+      .map((s) => s.mcp_server_id as import('@agor/core/types').MCPServerID);
+
+    if (serverIds.length > 0) {
+      const batchedTokens = await userTokenRepo.getTokensForServers(
+        userId as import('@agor/core/types').UserID,
+        serverIds
+      );
+
+      for (const token of batchedTokens) {
+        const keyUserId = token.user_id ?? 'shared';
+        tokenMap.set(`${token.mcp_server_id}:${keyUserId}`, token);
+      }
+    }
+
     const injectToken = async (server: MCPServer) => {
       if (server.auth?.type !== 'oauth') {
         return server;
@@ -811,8 +841,8 @@ export function registerHooks(ctx: RegisterHooksContext): void {
         mode === 'per_user' ? (userId as import('@agor/core/types').UserID) : null;
 
       try {
-        const userTokenRepo = new UserMCPOAuthTokenRepository(db);
-        const row = await userTokenRepo.getToken(tokenUserId, server.mcp_server_id);
+        const keyUserId = tokenUserId ?? 'shared';
+        const row = tokenMap.get(`${server.mcp_server_id}:${keyUserId}`);
 
         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, or } from 'drizzle-orm';
 import type { Database } from '../client';
 import { deleteFrom, insert, select, update } from '../database-wrapper';
 import {
@@ -260,4 +260,30 @@ export class UserMCPOAuthTokenRepository {
     const token = await this.getValidToken(userId, serverId);
     return token !== undefined;
   }
+
+  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
+              ? or(isNull(userMcpOauthTokens.user_id), eq(userMcpOauthTokens.user_id, userId))
+              : isNull(userMcpOauthTokens.user_id)
+          )
+        )
+        .all();
+      return rows.map(rowToToken);
+    } catch (error) {
+      throw new RepositoryError(
+        `Failed to get OAuth tokens for servers: ${error instanceof Error ? error.message : String(error)}`,
+        error
+      );
+    }
+  }
 }