From 7c6c30cb6ff26bbecdc5ef1b70e185fcdece9b1a Mon Sep 17 00:00:00 2001 From: Donach <39565367+Donach@users.noreply.github.com> Date: Mon, 20 Jul 2026 07:28:13 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Batch=20MCP=20OAuth=20token?= =?UTF-8?q?=20fetches=20to=20fix=20N+1=20query?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .jules/bolt.md | 3 ++ apps/agor-daemon/src/register-hooks.ts | 49 ++++++++++++++++++- .../db/repositories/user-mcp-oauth-tokens.ts | 31 +++++++++++- 3 files changed, 80 insertions(+), 3 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 8c4087435e..3139a434b0 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. +## 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., `:` and `:shared`), and look up the data from this Map during the loop. diff --git a/apps/agor-daemon/src/register-hooks.ts b/apps/agor-daemon/src/register-hooks.ts index 05321d3669..6d4cdfeb93 100755 --- a/apps/agor-daemon/src/register-hooks.ts +++ b/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(); + 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( 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..a2ac8d2662 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 } 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 { + 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.