From 3094579a2cd292208bdda93dab53f81cf7ce690d Mon Sep 17 00:00:00 2001 From: Donach <39565367+Donach@users.noreply.github.com> Date: Thu, 30 Jul 2026 07:11:40 +0000 Subject: [PATCH] perf: batch MCP OAuth token queries to fix N+1 Batch resolves MCP OAuth tokens using an in-memory Map initialized via a single `getTokensForServers` database query to eliminate N+1 latency in `mcp-servers` and `session-mcp-servers` API endpoints. --- .jules/bolt.md | 3 ++ apps/agor-daemon/src/register-hooks.ts | 36 +++++++++++++++++-- .../db/repositories/user-mcp-oauth-tokens.ts | 28 ++++++++++++++- 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 8c4087435e..9727932b27 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. +## 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. `:`) 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. diff --git a/apps/agor-daemon/src/register-hooks.ts b/apps/agor-daemon/src/register-hooks.ts index 05321d3669..b89516d174 100755 --- a/apps/agor-daemon/src/register-hooks.ts +++ b/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(); + + 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( 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..7a5e0515f8 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 { @@ -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 { + 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 + ); + } + } }