From 86fb8f0d5ba7dbdba4e39b96405fa65ad32e9791 Mon Sep 17 00:00:00 2001 From: Donach <39565367+Donach@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:10:41 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Pre-fetch=20OAuth=20tokens?= =?UTF-8?q?=20to=20eliminate=20N+1=20query?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Added `getTokensForServers` batched lookup to `UserMCPOAuthTokenRepository`. * Refactored `injectPerUserOAuthTokens` hook in `register-hooks.ts` to fetch all necessary tokens upfront and map them by `serverId:userId`. * Avoids repeated database roundtrips when loading multiple MCP servers. --- .jules/bolt.md | 3 ++ apps/agor-daemon/src/register-hooks.ts | 32 ++++++++++++++++-- .../db/repositories/user-mcp-oauth-tokens.ts | 33 ++++++++++++++++++- 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 8c4087435e..281cbfb270 100644 --- a/.jules/bolt.md +++ b/.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. diff --git a/apps/agor-daemon/src/register-hooks.ts b/apps/agor-daemon/src/register-hooks.ts index 05321d3669..48b583d63c 100755 --- a/apps/agor-daemon/src/register-hooks.ts +++ b/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) { diff --git a/packages/core/src/db/repositories/user-mcp-oauth-tokens.ts b/packages/core/src/db/repositories/user-mcp-oauth-tokens.ts index c5ea01f893..1f28d62ac2 100644 --- a/packages/core/src/db/repositories/user-mcp-oauth-tokens.ts +++ b/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 { + 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 { try { const row = await select(this.db)