From 3a3e6b55f58e10c908105b97134f69ecb47c1978 Mon Sep 17 00:00:00 2001 From: Donach <39565367+Donach@users.noreply.github.com> Date: Sun, 2 Aug 2026 07:19:37 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Batch=20MCP=20OAuth=20token?= =?UTF-8?q?=20injection=20to=20fix=20N+1=20query=20in=20FeathersJS=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 💡 What: Replaced individual database queries for OAuth tokens inside a `Promise.all` FeathersJS map with a single batched database query using `inArray` and an in-memory Map lookup. 🎯 Why: Resolves a critical N+1 database querying bottleneck during the `injectPerUserOAuthTokens` hook whenever an array of MCP servers is fetched. 📊 Impact: Reduces O(N) database roundtrips for token lookups down to a single O(1) query per hook invocation, significantly speeding up server listing and serialization logic for large boards. 🔬 Measurement: Verified with `pnpm test` ensuring existing auth scoping logic holds. Performance improvement will be visible in reduced latency during `mcp-servers` service listing operations. --- .jules/bolt.md | 3 + apps/agor-daemon/src/register-hooks.ts | 155 +++++++++--------- .../db/repositories/user-mcp-oauth-tokens.ts | 32 +++- 3 files changed, 116 insertions(+), 74 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 8c4087435e..6a9c7aef5a 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. +## $(date +%Y-%m-%d) - [Batch FeathersJS token injection with inArray to fix N+1 query] +**Learning:** When retrieving records during a FeathersJS hook (e.g. `after: { find: [injectPerUserOAuthTokens] }`), making a database lookup per item inside `Promise.all(context.result.map(injectToken))` results in an N+1 query problem. Pre-fetching all necessary relationships using a single `inArray()` lookup based on a map eliminates this performance bottleneck. +**Action:** Always extract all relevant foreign keys from the `context.result` array, perform a single batched database lookup, construct a mapping, and then sync the array with the map in memory instead of executing a query per record. diff --git a/apps/agor-daemon/src/register-hooks.ts b/apps/agor-daemon/src/register-hooks.ts index 05321d3669..9369ced466 100755 --- a/apps/agor-daemon/src/register-hooks.ts +++ b/apps/agor-daemon/src/register-hooks.ts @@ -798,90 +798,99 @@ export function registerHooks(ctx: RegisterHooksContext): void { return context; } - const injectToken = async (server: MCPServer) => { - if (server.auth?.type !== 'oauth') { - return server; - } + // Batch fetching of OAuth tokens to eliminate N+1 queries + const servers = Array.isArray(context.result) + ? context.result + : context.result?.data && Array.isArray(context.result.data) + ? context.result.data + : context.result?.mcp_server_id + ? [context.result] + : []; + + if (servers.length > 0) { + const oauthServers = servers.filter((s: MCPServer) => s.auth?.type === 'oauth'); + if (oauthServers.length > 0) { + const userTokenRepo = new UserMCPOAuthTokenRepository(db); + const serverIds = oauthServers.map((s: MCPServer) => s.mcp_server_id); + const tokenMap = await userTokenRepo.getTokensForServers( + userId as import('@agor/core/types').UserID, + serverIds + ); - // Tokens for both modes live in user_mcp_oauth_tokens: - // - per_user → row keyed by (userId, serverId) - // - shared → row keyed by (NULL, serverId) - const mode = server.auth.oauth_mode ?? 'per_user'; - const tokenUserId: import('@agor/core/types').UserID | null = - mode === 'per_user' ? (userId as import('@agor/core/types').UserID) : null; + // Pass map to a modified injectToken that takes the pre-fetched map + const batchInjectToken = async (server: MCPServer) => { + if (server.auth?.type !== 'oauth') return server; - try { - const userTokenRepo = new UserMCPOAuthTokenRepository(db); - const row = await userTokenRepo.getToken(tokenUserId, server.mcp_server_id); + const mode = server.auth.oauth_mode ?? 'per_user'; + const tokenUserId = + mode === 'per_user' ? (userId as import('@agor/core/types').UserID) : null; + const key = `${server.mcp_server_id}:${tokenUserId ?? 'shared'}`; + const row = tokenMap.get(key); - if (!row) { - console.log( - `[MCP OAuth] No token row for user=${tokenUserId ?? ''} server=${server.name}` - ); - return server; - } + if (!row) { + console.log( + `[MCP OAuth] No token row for user=${tokenUserId ?? ''} server=${server.name}` + ); + return server; + } - // JIT refresh — see `refreshAndPersistToken` for mutexing + invalid_grant cleanup. - let accessToken = row.oauth_access_token; - let expiresAt = row.oauth_token_expires_at; - const { needsRefresh, refreshAndPersistToken, InvalidGrantError } = await import( - '@agor/core/tools/mcp/oauth-refresh' - ); - if (needsRefresh(row.oauth_token_expires_at) && row.oauth_refresh_token) { - console.log(`[MCP OAuth] Token near/past expiry for ${server.name} — refreshing`); try { - accessToken = await refreshAndPersistToken({ - db, - userId: tokenUserId, - mcpServerId: server.mcp_server_id, - }); - // Re-read to pick up the rotated expiry for the UI. - const fresh = await userTokenRepo.getToken(tokenUserId, server.mcp_server_id); - if (fresh) expiresAt = fresh.oauth_token_expires_at; - } catch (refreshErr) { - if (refreshErr instanceof InvalidGrantError) { - console.warn( - `[MCP OAuth] invalid_grant refreshing ${server.name} — user must re-auth` - ); - return server; + // JIT refresh + let accessToken = row.oauth_access_token; + let expiresAt = row.oauth_token_expires_at; + const { needsRefresh, refreshAndPersistToken, InvalidGrantError } = await import( + '@agor/core/tools/mcp/oauth-refresh' + ); + if (needsRefresh(row.oauth_token_expires_at) && row.oauth_refresh_token) { + console.log(`[MCP OAuth] Token near/past expiry for ${server.name} — refreshing`); + try { + accessToken = await refreshAndPersistToken({ + db, + userId: tokenUserId, + mcpServerId: server.mcp_server_id, + }); + const fresh = await userTokenRepo.getToken(tokenUserId, server.mcp_server_id); + if (fresh) expiresAt = fresh.oauth_token_expires_at; + } catch (refreshErr) { + if (refreshErr instanceof InvalidGrantError) { + console.warn( + `[MCP OAuth] invalid_grant refreshing ${server.name} — user must re-auth` + ); + return server; + } + console.warn( + `[MCP OAuth] Refresh failed for ${server.name} (using stale token):`, + refreshErr instanceof Error ? refreshErr.message : refreshErr + ); + } } - // Transient error: fall through with the stale access_token. The - // MCP call may still succeed or fail cleanly at the transport. + + return { + ...server, + auth: { + ...server.auth, + oauth_access_token: accessToken, + oauth_token_expires_at: + expiresAt instanceof Date ? expiresAt.getTime() : (expiresAt ?? undefined), + }, + }; + } catch (error) { console.warn( - `[MCP OAuth] Refresh failed for ${server.name} (using stale token):`, - refreshErr instanceof Error ? refreshErr.message : refreshErr + `[MCP OAuth] Failed to resolve OAuth token for ${server.name}:`, + error instanceof Error ? error.message : error ); } - } - - return { - ...server, - auth: { - ...server.auth, - oauth_access_token: accessToken, - // Surface expiry so the UI can render "expires in X" tooltips. - // Stored as Date in the repo, emitted as ms epoch to match MCPAuth. - oauth_token_expires_at: - expiresAt instanceof Date ? expiresAt.getTime() : (expiresAt ?? undefined), - }, + return server; }; - } catch (error) { - console.warn( - `[MCP OAuth] Failed to resolve OAuth token for ${server.name}:`, - error instanceof Error ? error.message : error - ); - } - return server; - }; - - // Handle both single result and array/paginated results - if (Array.isArray(context.result)) { - context.result = await Promise.all(context.result.map(injectToken)); - } else if (context.result?.data && Array.isArray(context.result.data)) { - context.result.data = await Promise.all(context.result.data.map(injectToken)); - } else if (context.result?.mcp_server_id) { - context.result = await injectToken(context.result); + if (Array.isArray(context.result)) { + context.result = await Promise.all(context.result.map(batchInjectToken)); + } else if (context.result?.data && Array.isArray(context.result.data)) { + context.result.data = await Promise.all(context.result.data.map(batchInjectToken)); + } else if (context.result?.mcp_server_id) { + context.result = await batchInjectToken(context.result); + } + } } return context; 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..bba62042fd 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 { @@ -76,6 +76,36 @@ function matchKey(userId: UserID | null, serverId: MCPServerID) { export class UserMCPOAuthTokenRepository { constructor(private db: Database) {} + async getTokensForServers( + userId: UserID | null, + serverIds: MCPServerID[] + ): Promise> { + try { + if (serverIds.length === 0) return new Map(); + const conditions = []; + if (userId) conditions.push(eq(userMcpOauthTokens.user_id, userId)); + conditions.push(isNull(userMcpOauthTokens.user_id)); + + const rows = await select(this.db) + .from(userMcpOauthTokens) + .where(and(inArray(userMcpOauthTokens.mcp_server_id, serverIds), or(...conditions))) + .all(); + + const map = new Map(); + for (const row of rows) { + const token = rowToToken(row); + const key = `${token.mcp_server_id}:${token.user_id ?? 'shared'}`; + map.set(key, token); + } + return map; + } catch (error) { + throw new RepositoryError( + `Failed to get OAuth tokens for servers: ${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.