diff --git a/.jules/bolt.md b/.jules/bolt.md index 8c4087435e..8d74a3730c 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-02-12 - [Batch MCP OAuth tokens with compound keys] +**Learning:** When batching MCP OAuth token queries across multiple servers, the records can contain a mix of per-user tokens (user_id set) and shared tokens (user_id is null). Squashing these into a Map keyed only by `mcp_server_id` overwrites records and causes the wrong token to be injected based on the server's `oauth_mode`. +**Action:** When populating an in-memory Map from a batched database result for MCP tokens, always use a compound key like `:` to preserve both shared and per-user records, allowing the caller to select the correct one. diff --git a/apps/agor-daemon/src/register-hooks.ts b/apps/agor-daemon/src/register-hooks.ts index 05321d3669..304a14e399 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 in the result to eliminate N+1 queries. + const items = 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 mcpServerIds = items + .filter((s: MCPServer) => s.auth?.type === 'oauth') + .map((s: MCPServer) => s.mcp_server_id); + + const tokensMap = new Map(); + const userTokenRepo = new UserMCPOAuthTokenRepository(db); + + if (mcpServerIds.length > 0) { + const tokens = await userTokenRepo.getTokensForServers( + userId as import('@agor/core/types').UserID, + mcpServerIds + ); + // Store using a compound key to preserve both shared and per-user tokens + for (const token of tokens) { + tokensMap.set(`${token.mcp_server_id}:${token.user_id ?? 'shared'}`, 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 tokenKey = `${server.mcp_server_id}:${tokenUserId ?? 'shared'}`; + const row = tokensMap.get(tokenKey); 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..e90b300165 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 { @@ -240,6 +240,36 @@ export class UserMCPOAuthTokenRepository { } } + /** + * Fetch tokens for a list of servers, resolving both per-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), + userId === null + ? isNull(userMcpOauthTokens.user_id) + : or(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 + ); + } + } + async listForUser(userId: UserID): Promise { try { const rows = await select(this.db)