fix(mcp): cache result of discoverServerTools to prevent post-OAuth refetch storm (#4701)

* fix(mcp): write per-server cache from discoverServerTools to prevent post-OAuth refetch storm

* improvement(mcp): update server status from discoverServerTools + cap list-tools timeout at 30s

Author: waleedlatif1

apps/sim/app/api/mcp/oauth/callback/route.ts

@@ -167,7 +167,8 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
     }
 
     try {
-      await mcpService.clearCache(server.workspaceId)
+      // discoverServerTools writes the result to this server's cache so the UI's
+      // immediate refetch hits it instead of re-fetching live.
       await mcpService.discoverServerTools(session.user.id, server.id, server.workspaceId)
     } catch (e) {
       logger.warn('Post-auth tools refresh failed', toError(e).message)

apps/sim/lib/mcp/client.test.ts

@@ -37,6 +37,7 @@ vi.mock('@modelcontextprotocol/sdk/types.js', () => ({
 
 vi.mock('@/lib/core/execution-limits', () => ({
   getMaxExecutionTimeout: vi.fn().mockReturnValue(30000),
+  DEFAULT_EXECUTION_TIMEOUT_MS: 30000,
 }))
 
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'

apps/sim/lib/mcp/client.ts

@@ -36,6 +36,7 @@ import {
   type McpToolsChangedCallback,
   type McpVersionInfo,
 } from '@/lib/mcp/types'
+import { MCP_CLIENT_CONSTANTS } from '@/lib/mcp/utils'
 
 const logger = createLogger('McpClient')
 
@@ -167,7 +168,9 @@ export class McpClient {
     }
 
     try {
-      const result: ListToolsResult = await this.client.listTools()
+      const result: ListToolsResult = await this.client.listTools(undefined, {
+        timeout: MCP_CLIENT_CONSTANTS.LIST_TOOLS_TIMEOUT_MS,
+      })
 
       if (!result.tools || !Array.isArray(result.tools)) {
         logger.warn(`Invalid tools response from server ${this.config.name}:`, result)

apps/sim/lib/mcp/service.test.ts

@@ -38,13 +38,20 @@ const {
 })
 
 vi.mock('@sim/db', () => {
+  // `where(...)` resolves to the workspace's rows AND exposes `.limit()` for
+  // chains like `getServerConfig` that do `select().from().where().limit(1)`.
+  const where = (...args: unknown[]) => {
+    const rowsPromise = Promise.resolve(mockGetWorkspaceServersRows(...args))
+    const thenable = Object.assign(rowsPromise, {
+      limit: (n: number) => rowsPromise.then((rows) => rows.slice(0, n)),
+    })
+    return thenable
+  }
   const setter = vi.fn().mockReturnValue({ where: vi.fn().mockResolvedValue(undefined) })
   return {
     db: {
       select: vi.fn().mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: (...args: unknown[]) => mockGetWorkspaceServersRows(...args),
-        }),
+        from: vi.fn().mockReturnValue({ where }),
       }),
       update: vi.fn().mockReturnValue({ set: setter }),
       insert: vi.fn(),
@@ -238,4 +245,18 @@ describe('McpService.discoverTools per-server caching', () => {
     expect(second.map((t) => t.name)).toEqual(['a-other'])
     expect(mockListTools).toHaveBeenCalledTimes(2)
   })
+
+  it('discoverServerTools primes the per-server cache for follow-up discoverTools', async () => {
+    mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')])
+    mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')])
+
+    const tools = await mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID)
+    expect(tools.map((t) => t.name)).toEqual(['a1'])
+    expect(mockListTools).toHaveBeenCalledTimes(1)
+
+    mockListTools.mockClear()
+    const second = await mcpService.discoverTools(USER_ID, WORKSPACE_ID)
+    expect(second.map((t) => t.name)).toEqual(['a1'])
+    expect(mockListTools).not.toHaveBeenCalled()
+  })
 })

apps/sim/lib/mcp/service.ts

@@ -580,6 +580,17 @@ class McpService {
         try {
           const tools = await client.listTools()
           logger.info(`[${requestId}] Discovered ${tools.length} tools from server ${config.name}`)
+          // Prime the per-server cache and reflect the successful connection on
+          // the row so the UI doesn't keep showing "Connect with OAuth" or stale
+          // disconnected/error state.
+          await Promise.allSettled([
+            this.cacheAdapter
+              .set(serverCacheKey(workspaceId, serverId), tools, this.cacheTimeout)
+              .catch((err) =>
+                logger.warn(`[${requestId}] Cache write failed for ${config.name}:`, err)
+              ),
+            this.updateServerStatus(serverId, workspaceId, true, undefined, tools.length),
+          ])
           return tools
         } finally {
           await client.disconnect()

apps/sim/lib/mcp/utils.ts

@@ -46,6 +46,8 @@ export function sanitizeHeaders(
 export const MCP_CLIENT_CONSTANTS = {
   CLIENT_TIMEOUT: DEFAULT_EXECUTION_TIMEOUT_MS,
   AUTO_REFRESH_INTERVAL: 5 * 60 * 1000,
+  // Cap metadata calls so a slow upstream can't hang the UI for 60s+.
+  LIST_TOOLS_TIMEOUT_MS: 30_000,
 } as const
 
 /**