⚡ Bolt: [Performance] Eliminate N+1 query when fetching MCP OAuth tokens

.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.
+## 2023-11-20 - [Fix N+1 query fetching OAuth tokens for servers in Feathers hooks]
+**Learning:** When enriching a list of records in a FeathersJS `after` hook (like mapping over `context.result` to attach OAuth tokens to MCP servers), fetching records inside a `Promise.all` loop causes an N+1 query problem. Pre-fetching them in a single query and loading them into an in-memory Map avoids excessive round-trips. When the key for the Map involves a nullable column (e.g. `user_id` being null for shared tokens), represent this safely in the Map's string key (e.g., using `'shared'`) to prevent collisions or undefined behavior. Furthermore, when crossing monorepo workspace boundaries, import types using the package's public export alias (e.g., `@agor/core/db`) instead of deep-linking internal files (e.g., `@agor/core/db/repositories/...`) to prevent module resolution failures during `pnpm build`.
+**Action:** When implementing or updating Feathers hooks that enrich arrays of results, always collect the necessary IDs, execute a single batched lookup query, and build an in-memory Map for O(1) assignment. Be careful to ensure the correct package export aliases are used for types.

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

@@ -798,6 +798,33 @@ export function registerHooks(ctx: RegisterHooksContext): void {
       return context;
     }
 
+    // Pre-fetch tokens for all servers to avoid N+1 query issue
+    const serverList: MCPServer[] = Array.isArray(context.result)
+      ? context.result
+      : context.result?.data && Array.isArray(context.result.data)
+        ? context.result.data
+        : context.result?.mcp_server_id
+          ? [context.result]
+          : [];
+
+    const serverIds = serverList
+      .filter((s) => s.auth?.type === 'oauth')
+      .map((s) => s.mcp_server_id);
+
+    const tokenMap = new Map<
+      string,
+      import('@agor/core/db').UserMCPOAuthToken
+    >();
+    if (serverIds.length > 0) {
+      const userTokenRepo = new UserMCPOAuthTokenRepository(db);
+      const typedUserId = userId as import('@agor/core/types').UserID;
+      const tokens = await userTokenRepo.getTokensForServers(typedUserId, serverIds);
+      for (const token of tokens) {
+        const key = `${token.mcp_server_id}:${token.user_id === null ? 'shared' : token.user_id}`;
+        tokenMap.set(key, token);
+      }
+    }
+
     const injectToken = async (server: MCPServer) => {
       if (server.auth?.type !== 'oauth') {
         return server;
@@ -811,8 +838,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 key = `${server.mcp_server_id}:${tokenUserId === null ? 'shared' : tokenUserId}`;
+        const row = tokenMap.get(key);
 
         if (!row) {
           console.log(
@@ -836,6 +863,7 @@ export function registerHooks(ctx: RegisterHooksContext): void {
               mcpServerId: server.mcp_server_id,
             });
             // Re-read to pick up the rotated expiry for the UI.
+            const userTokenRepo = new UserMCPOAuthTokenRepository(db);
             const fresh = await userTokenRepo.getToken(tokenUserId, server.mcp_server_id);
             if (fresh) expiresAt = fresh.oauth_token_expires_at;
           } catch (refreshErr) {

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 {
@@ -80,6 +80,37 @@ export class UserMCPOAuthTokenRepository {
    * Look up the token row for a (user, server) pair. Pass `null` for userId
    * to read the shared-mode row.
    */
+  /**
+   * Batch look up token rows for a set of servers and a specific user.
+   * Includes both per-user tokens for the user and shared-mode tokens.
+   */
+  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),
+            or(
+              isNull(userMcpOauthTokens.user_id),
+              userId ? eq(userMcpOauthTokens.user_id, userId) : undefined
+            )
+          )
+        )
+        .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
+      );
+    }
+  }
+
   async getToken(userId: UserID | null, serverId: MCPServerID): Promise<UserMCPOAuthToken | null> {
     try {
       const row = await select(this.db)