From 8412ea581d8ccece1806c122b10175a05e74c646 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 21 May 2026 10:46:39 -0700 Subject: [PATCH 1/2] fix(mcp): write per-server cache from discoverServerTools to prevent post-OAuth refetch storm --- apps/sim/app/api/mcp/oauth/callback/route.ts | 3 ++- apps/sim/lib/mcp/service.test.ts | 27 +++++++++++++++++--- apps/sim/lib/mcp/service.ts | 7 +++++ 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/apps/sim/app/api/mcp/oauth/callback/route.ts b/apps/sim/app/api/mcp/oauth/callback/route.ts index 721675dac9..15b115c0b7 100644 --- a/apps/sim/app/api/mcp/oauth/callback/route.ts +++ b/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) diff --git a/apps/sim/lib/mcp/service.test.ts b/apps/sim/lib/mcp/service.test.ts index 539ee38457..42d111640d 100644 --- a/apps/sim/lib/mcp/service.test.ts +++ b/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() + }) }) diff --git a/apps/sim/lib/mcp/service.ts b/apps/sim/lib/mcp/service.ts index f3ca6683ae..51ec434182 100644 --- a/apps/sim/lib/mcp/service.ts +++ b/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() From 4e6210c5cb7b6b5612d3dedcc3f5734af34ce8e2 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 21 May 2026 10:55:36 -0700 Subject: [PATCH 2/2] improvement(mcp): update server status from discoverServerTools + cap list-tools timeout at 30s --- apps/sim/lib/mcp/client.test.ts | 1 + apps/sim/lib/mcp/client.ts | 5 ++++- apps/sim/lib/mcp/service.ts | 18 +++++++++++------- apps/sim/lib/mcp/utils.ts | 2 ++ 4 files changed, 18 insertions(+), 8 deletions(-) diff --git a/apps/sim/lib/mcp/client.test.ts b/apps/sim/lib/mcp/client.test.ts index 405b66610e..8f6279a2d1 100644 --- a/apps/sim/lib/mcp/client.test.ts +++ b/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' diff --git a/apps/sim/lib/mcp/client.ts b/apps/sim/lib/mcp/client.ts index 819afc5660..9f3c36d00a 100644 --- a/apps/sim/lib/mcp/client.ts +++ b/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) diff --git a/apps/sim/lib/mcp/service.ts b/apps/sim/lib/mcp/service.ts index 51ec434182..b024c17930 100644 --- a/apps/sim/lib/mcp/service.ts +++ b/apps/sim/lib/mcp/service.ts @@ -580,13 +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 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) - ) + // 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() diff --git a/apps/sim/lib/mcp/utils.ts b/apps/sim/lib/mcp/utils.ts index 1fc231c7c6..af16621567 100644 --- a/apps/sim/lib/mcp/utils.ts +++ b/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 /**