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

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/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,13 @@ class McpService {
         try {
           const tools = await client.listTools()
           logger.info(`[${requestId}] Discovered ${tools.length} tools from server ${config.name}`)
+          // Prime the per-server cache so the next discoverTools call hits it
+          // instead of re-fetching live (which can stall on slow upstreams).
+          await this.cacheAdapter
+            .set(serverCacheKey(workspaceId, serverId), tools, this.cacheTimeout)
+            .catch((err) =>
+              logger.warn(`[${requestId}] Cache write failed for ${config.name}:`, err)
+            )
           return tools
         } finally {
           await client.disconnect()