From abaad6e9a11a98e0295e76998e50b53bb120b5ec Mon Sep 17 00:00:00 2001
From: Donach <39565367+Donach@users.noreply.github.com>
Date: Tue, 4 Aug 2026 07:19:41 +0000
Subject: [PATCH] perf(daemon): batch MCP OAuth token database lookups

Replaces the N+1 database queries for MCP server OAuth tokens
with a single batched query using Drizzle's `inArray` inside
`apps/agor-daemon/src/register-hooks.ts`.

Tokens are now pre-fetched and stored in an O(1) in-memory Map
using a compound key of `<server_id>:<user_id|shared>`, preserving
both shared-mode and per-user token isolation.
---
 apps/agor-daemon/src/register-hooks.ts        | 32 +++++++++++++++-
 .../db/repositories/user-mcp-oauth-tokens.ts  | 37 ++++++++++++++++++-
 2 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/apps/agor-daemon/src/register-hooks.ts b/apps/agor-daemon/src/register-hooks.ts
index 05321d3669..7da8adb2a6 100755
--- a/apps/agor-daemon/src/register-hooks.ts
+++ b/apps/agor-daemon/src/register-hooks.ts
@@ -798,6 +798,34 @@ export function registerHooks(ctx: RegisterHooksContext): void {
       return context;
     }
 
+    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];
+    }
+
+    const oauthServers = servers.filter((s) => s.auth?.type === 'oauth');
+    const userTokenRepo = new UserMCPOAuthTokenRepository(db);
+    const tokenMap = new Map<string, import('@agor/core/db').UserMCPOAuthToken>();
+
+    if (oauthServers.length > 0) {
+      const serverIds = oauthServers.map(
+        (s) => s.mcp_server_id as import('@agor/core/types').MCPServerID
+      );
+      const batchedTokens = await userTokenRepo.getTokensForServers(
+        serverIds,
+        userId as import('@agor/core/types').UserID
+      );
+
+      for (const token of batchedTokens) {
+        const key = `${token.mcp_server_id}:${token.user_id ?? 'shared'}`;
+        tokenMap.set(key, token);
+      }
+    }
+
     const injectToken = async (server: MCPServer) => {
       if (server.auth?.type !== 'oauth') {
         return server;
@@ -811,8 +839,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 ?? 'shared'}`;
+        const row = tokenMap.get(key);
 
         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..82a595c8e7 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,41 @@ export class UserMCPOAuthTokenRepository {
    * Look up the token row for a (user, server) pair. Pass `null` for userId
    * to read the shared-mode row.
    */
+  /**
+   * Look up tokens for a list of servers. Fetches both shared-mode and per-user
+   * tokens for the given userId.
+   */
+  async getTokensForServers(
+    serverIds: MCPServerID[],
+    userId: UserID | null
+  ): Promise<UserMCPOAuthToken[]> {
+    if (serverIds.length === 0) return [];
+
+    try {
+      const conditions: any[] = [inArray(userMcpOauthTokens.mcp_server_id, serverIds)];
+
+      if (userId) {
+        conditions.push(
+          or(isNull(userMcpOauthTokens.user_id), eq(userMcpOauthTokens.user_id, userId))
+        );
+      } else {
+        conditions.push(isNull(userMcpOauthTokens.user_id));
+      }
+
+      const rows = await select(this.db)
+        .from(userMcpOauthTokens)
+        .where(and(...conditions))
+        .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<UserMCPOAuthToken | null> {
     try {
       const row = await select(this.db)