diff --git a/apps/sim/app/api/mcp/oauth/start/route.test.ts b/apps/sim/app/api/mcp/oauth/start/route.test.ts index 68dd96b2bf3..555e2d2fa4d 100644 --- a/apps/sim/app/api/mcp/oauth/start/route.test.ts +++ b/apps/sim/app/api/mcp/oauth/start/route.test.ts @@ -162,6 +162,23 @@ describe('MCP OAuth start route', () => { expect(body.error).not.toContain('ECONNREFUSED') expect(body.error).not.toContain('internal-db-host') }) + + it('returns an actionable 4xx when the server does not support dynamic client registration', async () => { + mcpOauthMockFns.mockMcpAuthGuarded.mockRejectedValueOnce( + new Error('Incompatible auth server: does not support dynamic client registration') + ) + const request = new NextRequest( + 'http://localhost:3000/api/mcp/oauth/start?workspaceId=workspace-1&serverId=server-1' + ) + + const response = await GET(request) + const body = await response.json() + + expect(response.status).toBe(422) + expect(body.error).toBe( + "This server doesn't support OAuth client registration. Configure a token instead." + ) + }) }) describe('surfaceOauthError', () => { diff --git a/apps/sim/app/api/mcp/oauth/start/route.ts b/apps/sim/app/api/mcp/oauth/start/route.ts index edb17a0f1c9..4354257a1e0 100644 --- a/apps/sim/app/api/mcp/oauth/start/route.ts +++ b/apps/sim/app/api/mcp/oauth/start/route.ts @@ -2,7 +2,7 @@ import { OAuthError, ServerError } from '@modelcontextprotocol/sdk/server/auth/e import { db } from '@sim/db' import { mcpServers } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { toError } from '@sim/utils/errors' +import { getErrorMessage, toError } from '@sim/utils/errors' import { and, eq, isNull } from 'drizzle-orm' import type { NextRequest } from 'next/server' import { NextResponse } from 'next/server' @@ -25,6 +25,14 @@ import { createMcpErrorResponse } from '@/lib/mcp/utils' const logger = createLogger('McpOauthStartAPI') const OAUTH_START_TTL_MS = 10 * 60 * 1000 const MAX_SURFACED_ERROR_LENGTH = 250 +const DCR_UNSUPPORTED_MESSAGE = + "This server doesn't support OAuth client registration. Configure a token instead." + +function isDynamicClientRegistrationUnsupported(error: unknown): boolean { + return getErrorMessage(error, '') + .toLowerCase() + .includes('does not support dynamic client registration') +} export function surfaceOauthError(error: unknown): string { // Spec-compliant OAuth servers throw typed subclasses with clean RFC 6749 fields. @@ -148,6 +156,9 @@ export const GET = withRouteHandler( authorizationUrl: e.authorizationUrl, }) } + if (isDynamicClientRegistrationUnsupported(e)) { + return createMcpErrorResponse(toError(e), DCR_UNSUPPORTED_MESSAGE, 422) + } throw e } } catch (error) { diff --git a/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts b/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts index e3f217fbdde..62950741392 100644 --- a/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts +++ b/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts @@ -3,18 +3,21 @@ */ import type { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' +import { McpOauthAuthorizationRequiredError } from '@/lib/mcp/types' -const { mockClearCache, mockDiscoverServerTools, mockSelect, mockUpdateSet } = vi.hoisted(() => ({ - mockClearCache: vi.fn(), - mockDiscoverServerTools: vi.fn(), - mockSelect: vi.fn(), - mockUpdateSet: vi.fn(), -})) +const { mockClearCache, mockDiscoverServerTools, mockSelect, mockUpdate, mockUpdateSet } = + vi.hoisted(() => ({ + mockClearCache: vi.fn(), + mockDiscoverServerTools: vi.fn(), + mockSelect: vi.fn(), + mockUpdate: vi.fn(), + mockUpdateSet: vi.fn(), + })) vi.mock('@sim/db', () => ({ db: { select: mockSelect, - update: vi.fn().mockReturnValue({ set: mockUpdateSet }), + update: mockUpdate.mockReturnValue({ set: mockUpdateSet }), }, })) @@ -43,7 +46,7 @@ vi.mock('@/lib/mcp/middleware', () => ({ vi.mock('@/lib/mcp/service', () => ({ mcpService: { clearCache: mockClearCache, - discoverServerTools: mockDiscoverServerTools, + discoverServerToolsWithMetadata: mockDiscoverServerTools, }, })) @@ -54,9 +57,11 @@ const initialServer = { workspaceId: 'workspace-1', name: 'OAuth Server', url: 'https://example.com/mcp', + authType: 'oauth', connectionStatus: 'connected', lastError: null, lastConnected: new Date('2026-01-01T00:00:00.000Z'), + lastToolsRefresh: new Date('2026-01-01T00:00:00.000Z'), toolCount: 4, statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, } @@ -81,6 +86,8 @@ function selectRows(rows: unknown[]) { describe('MCP server refresh route', () => { beforeEach(() => { vi.clearAllMocks() + mockSelect.mockReset() + mockSelect.mockReturnValue(selectRows([persistedServer])) mockSelect.mockReturnValueOnce(selectRows([initialServer])) mockUpdateSet.mockReturnValue({ where: vi.fn().mockReturnValue({ returning: vi.fn().mockResolvedValue([persistedServer]) }), @@ -88,7 +95,9 @@ describe('MCP server refresh route', () => { }) it('preserves the service-persisted OAuth pending status', async () => { - mockDiscoverServerTools.mockRejectedValueOnce(new Error('OAuth authorization required')) + mockDiscoverServerTools.mockRejectedValueOnce( + new McpOauthAuthorizationRequiredError('server-1', 'OAuth Server') + ) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', @@ -107,16 +116,53 @@ describe('MCP server refresh route', () => { ) }) + it('preserves OAuth pending when the status reread is still stale connected', async () => { + mockDiscoverServerTools.mockRejectedValueOnce( + new McpOauthAuthorizationRequiredError('server-1', 'OAuth Server') + ) + mockSelect.mockReturnValueOnce(selectRows([initialServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'disconnected', + error: null, + toolCount: 0, + workflowsUpdated: 0, + }) + ) + }) + + it('reports a sanitized discovery timeout when persistence leaves disconnected without an error', async () => { + mockDiscoverServerTools.mockRejectedValueOnce(new Error('upstream request timeout')) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'disconnected', + error: 'Request timed out', + toolCount: 0, + workflowsUpdated: 0, + }) + ) + }) + it('reports the discovery failure when status persistence leaves a stale connected row', async () => { const reflectedSecret = 'Bearer reflected-static-token' mockDiscoverServerTools.mockRejectedValueOnce( new Error(`Upstream reflected ${reflectedSecret}`) ) - mockUpdateSet.mockReturnValueOnce({ - where: vi.fn().mockReturnValue({ - returning: vi.fn().mockResolvedValue([initialServer]), - }), - }) + mockSelect.mockReturnValueOnce(selectRows([initialServer])) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', @@ -128,6 +174,7 @@ describe('MCP server refresh route', () => { expect.objectContaining({ status: 'disconnected', error: 'Internal server error', + toolCount: 0, workflowsUpdated: 0, }) ) @@ -140,13 +187,10 @@ describe('MCP server refresh route', () => { const newerSuccessfulServer = { ...initialServer, lastConnected: new Date(Date.now() + 60_000), + lastToolsRefresh: new Date(Date.now() + 60_000), toolCount: 7, } - mockUpdateSet.mockReturnValueOnce({ - where: vi.fn().mockReturnValue({ - returning: vi.fn().mockResolvedValue([newerSuccessfulServer]), - }), - }) + mockSelect.mockReturnValueOnce(selectRows([newerSuccessfulServer])) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', @@ -162,25 +206,52 @@ describe('MCP server refresh route', () => { workflowsUpdated: 0, }) ) - expect(mockClearCache).toHaveBeenCalledWith('workspace-1') + expect(mockClearCache).not.toHaveBeenCalled() + }) + + it('reports the discovery failure when only cache invalidation advanced', async () => { + mockDiscoverServerTools.mockRejectedValueOnce(new Error('Connection failed')) + const cacheInvalidatedServer = { + ...initialServer, + lastToolsRefresh: new Date(initialServer.lastToolsRefresh.getTime() + 60_000), + toolCount: 0, + } + mockSelect.mockReturnValueOnce(selectRows([cacheInvalidatedServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'disconnected', + error: 'Internal server error', + toolCount: 0, + workflowsUpdated: 0, + }) + ) + expect(mockClearCache).not.toHaveBeenCalled() }) it('does not 500 when workflow sync fails after a successful discovery', async () => { - mockDiscoverServerTools.mockResolvedValueOnce([ - { - name: 'search', - description: 'Search tool', - inputSchema: {}, - serverId: 'server-1', - serverName: 'OAuth Server', - }, - ]) + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [ + { + name: 'search', + description: 'Search tool', + inputSchema: {}, + serverId: 'server-1', + serverName: 'OAuth Server', + }, + ], + state: 'published', + }) // The route's server lookup consumes the first select (beforeEach). The sync's // workflow select is left unmocked, so it throws — exercising the guard that // keeps a secondary sync failure from turning a successful refresh into a 500. - mockUpdateSet.mockReturnValueOnce({ - where: vi.fn().mockReturnValue({ returning: vi.fn().mockResolvedValue([initialServer]) }), - }) + mockSelect.mockReturnValueOnce(selectRows([initialServer])) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', @@ -197,4 +268,210 @@ describe('MCP server refresh route', () => { }) ) }) + + it('reports live tools when cache degradation prevents status publication', async () => { + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [ + { + name: 'search', + description: 'Search tool', + inputSchema: {}, + serverId: 'server-1', + serverName: 'OAuth Server', + }, + ], + state: 'unavailable', + }) + mockSelect.mockReturnValueOnce(selectRows([persistedServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'connected', + error: null, + toolCount: 1, + }) + ) + }) + + it('returns the winning cached tools without syncing workflows from the losing refresh', async () => { + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [ + { + name: 'winner-search', + description: 'Search tool from the winning refresh', + inputSchema: {}, + serverId: 'server-1', + serverName: 'OAuth Server', + }, + ], + state: 'winner-cache', + }) + mockSelect.mockReturnValueOnce(selectRows([persistedServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'connected', + error: null, + toolCount: 1, + workflowsUpdated: 0, + updatedWorkflowIds: [], + }) + ) + expect(mockSelect).toHaveBeenCalledTimes(2) + expect(mockUpdate).not.toHaveBeenCalled() + }) + + it.each(['unavailable', 'winner-cache'] as const)( + 'preserves a newer cache invalidation over %s discovery tools', + async (state) => { + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [ + { + name: 'stale-search', + description: 'Search tool loaded before invalidation', + inputSchema: {}, + serverId: 'server-1', + serverName: 'OAuth Server', + }, + ], + state, + }) + const cacheInvalidatedServer = { + ...initialServer, + lastToolsRefresh: new Date(initialServer.lastToolsRefresh.getTime() + 60_000), + toolCount: 0, + } + mockSelect.mockReturnValueOnce(selectRows([cacheInvalidatedServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'connected', + error: null, + toolCount: 0, + workflowsUpdated: 0, + updatedWorkflowIds: [], + }) + ) + expect(mockSelect).toHaveBeenCalledTimes(2) + expect(mockUpdate).not.toHaveBeenCalled() + } + ) + + it('does not sync published tools after a newer cache invalidation', async () => { + const publicationOrder = new Date(initialServer.lastToolsRefresh.getTime() + 30_000) + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [ + { + name: 'stale-search', + description: 'Search tool published before invalidation', + inputSchema: {}, + serverId: 'server-1', + serverName: 'OAuth Server', + }, + ], + state: 'published', + publicationOrder, + }) + const cacheInvalidatedServer = { + ...initialServer, + lastConnected: publicationOrder, + lastToolsRefresh: new Date(publicationOrder.getTime() + 30_000), + toolCount: 0, + } + mockSelect.mockReturnValueOnce(selectRows([cacheInvalidatedServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'connected', + error: null, + toolCount: 0, + workflowsUpdated: 0, + updatedWorkflowIds: [], + }) + ) + expect(mockSelect).toHaveBeenCalledTimes(2) + expect(mockUpdate).not.toHaveBeenCalled() + }) + + it('preserves a newer successful refresh when discovery is superseded', async () => { + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [], + state: 'superseded', + }) + const newerSuccessfulServer = { + ...initialServer, + lastConnected: new Date(initialServer.lastConnected.getTime() + 60_000), + lastToolsRefresh: new Date(initialServer.lastToolsRefresh.getTime() + 60_000), + toolCount: 7, + } + mockSelect.mockReturnValueOnce(selectRows([newerSuccessfulServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'connected', + error: null, + toolCount: 7, + workflowsUpdated: 0, + updatedWorkflowIds: [], + }) + ) + expect(mockSelect).toHaveBeenCalledTimes(2) + expect(mockUpdate).not.toHaveBeenCalled() + }) + + it('fails closed without syncing workflows when discovery is superseded', async () => { + mockDiscoverServerTools.mockResolvedValueOnce({ + tools: [], + state: 'superseded', + }) + mockSelect.mockReturnValueOnce(selectRows([initialServer])) + + const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { + method: 'POST', + }) as NextRequest + const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) }) + const body = await response.json() + + expect(body.data).toEqual( + expect.objectContaining({ + status: 'disconnected', + error: 'Tool discovery was superseded by a newer refresh. Please retry.', + toolCount: 0, + workflowsUpdated: 0, + updatedWorkflowIds: [], + }) + ) + expect(mockSelect).toHaveBeenCalledTimes(2) + expect(mockUpdate).not.toHaveBeenCalled() + }) }) diff --git a/apps/sim/app/api/mcp/servers/[id]/refresh/route.ts b/apps/sim/app/api/mcp/servers/[id]/refresh/route.ts index 26bb6ddaaef..7cea24af0fd 100644 --- a/apps/sim/app/api/mcp/servers/[id]/refresh/route.ts +++ b/apps/sim/app/api/mcp/servers/[id]/refresh/route.ts @@ -1,3 +1,4 @@ +import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js' import { db } from '@sim/db' import { mcpServers, workflow, workflowBlocks } from '@sim/db/schema' import { createLogger } from '@sim/logger' @@ -9,8 +10,12 @@ import { mcpServerIdParamsSchema } from '@/lib/api/contracts/mcp' import { validationErrorResponse } from '@/lib/api/server' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' import { withMcpAuth } from '@/lib/mcp/middleware' -import { mcpService } from '@/lib/mcp/service' -import type { McpTool, McpToolSchema } from '@/lib/mcp/types' +import { type McpServerDiscoveryState, mcpService } from '@/lib/mcp/service' +import { + McpOauthAuthorizationRequiredError, + type McpTool, + type McpToolSchema, +} from '@/lib/mcp/types' import { categorizeError, createMcpErrorResponse, @@ -188,34 +193,71 @@ export const POST = withRouteHandler( let syncResult: SyncResult = { updatedCount: 0, updatedWorkflowIds: [] } let discoveredTools: McpTool[] = [] + let discoveryState: McpServerDiscoveryState | null = null + let discoveryPublicationOrder: Date | null = null let discoveryError: string | null = null - const discoveryStartedAt = new Date() + let oauthAuthorizationRequired = false try { - discoveredTools = await mcpService.discoverServerTools( + const discovery = await mcpService.discoverServerToolsWithMetadata( userId, serverId, workspaceId, true ) + discoveredTools = discovery.tools + discoveryState = discovery.state + discoveryPublicationOrder = discovery.publicationOrder ?? null logger.info( `[${requestId}] Discovered ${discoveredTools.length} tools from server ${serverId}` ) } catch (error) { + oauthAuthorizationRequired = + server.authType === 'oauth' && + (error instanceof McpOauthAuthorizationRequiredError || + error instanceof UnauthorizedError) discoveryError = truncate(categorizeError(error).message, 200, '') logger.warn(`[${requestId}] Failed to connect to server ${serverId}`, { error: discoveryError, }) } - if (discoveryError === null) { + const [refreshedServer] = await db + .select({ + name: mcpServers.name, + url: mcpServers.url, + connectionStatus: mcpServers.connectionStatus, + lastConnected: mcpServers.lastConnected, + lastToolsRefresh: mcpServers.lastToolsRefresh, + lastError: mcpServers.lastError, + toolCount: mcpServers.toolCount, + }) + .from(mcpServers) + .where( + and( + eq(mcpServers.id, serverId), + eq(mcpServers.workspaceId, workspaceId), + isNull(mcpServers.deletedAt) + ) + ) + .limit(1) + + const publicationBaseline = discoveryPublicationOrder ?? server.lastToolsRefresh + const newerPublicationWonRace = + refreshedServer?.lastToolsRefresh != null && + (publicationBaseline == null || refreshedServer.lastToolsRefresh > publicationBaseline) + + if (discoveryError === null && discoveryState === 'published' && !newerPublicationWonRace) { try { syncResult = await syncToolSchemasToWorkflows( workspaceId, serverId, discoveredTools, requestId, - { url: server.url ?? undefined, name: server.name ?? undefined } + { + url: refreshedServer?.url ?? undefined, + name: refreshedServer?.name ?? undefined, + } ) } catch (error) { // Discovery already persisted status and cached tools; a workflow-sync @@ -227,45 +269,43 @@ export const POST = withRouteHandler( } } - const now = new Date() - - const [refreshedServer] = await db - .update(mcpServers) - .set({ - lastToolsRefresh: now, - updatedAt: now, - }) - .where( - and( - eq(mcpServers.id, serverId), - eq(mcpServers.workspaceId, workspaceId), - isNull(mcpServers.deletedAt) - ) - ) - .returning({ - connectionStatus: mcpServers.connectionStatus, - lastConnected: mcpServers.lastConnected, - lastError: mcpServers.lastError, - toolCount: mcpServers.toolCount, - }) - let connectionStatus = refreshedServer?.connectionStatus ?? 'error' let lastError = refreshedServer ? refreshedServer.lastError : discoveryError - const toolCount = refreshedServer?.toolCount ?? discoveredTools.length + let toolCount = refreshedServer?.toolCount ?? discoveredTools.length + const newerSuccessWonRace = + connectionStatus === 'connected' && + newerPublicationWonRace && + refreshedServer.lastConnected != null && + (server.lastConnected == null || refreshedServer.lastConnected > server.lastConnected) + + if (discoveryState === 'superseded' && !newerSuccessWonRace) { + connectionStatus = 'disconnected' + lastError = 'Tool discovery was superseded by a newer refresh. Please retry.' + toolCount = 0 + } else if ( + discoveryError === null && + !newerPublicationWonRace && + (discoveryState === 'unavailable' || discoveryState === 'winner-cache') + ) { + connectionStatus = 'connected' + lastError = null + toolCount = discoveredTools.length + } if (discoveryError !== null && connectionStatus === 'connected') { - const newerSuccessWonRace = - refreshedServer?.lastConnected != null && - refreshedServer.lastConnected > discoveryStartedAt - if (!newerSuccessWonRace) { connectionStatus = 'disconnected' - lastError = discoveryError + lastError = oauthAuthorizationRequired ? null : discoveryError + toolCount = 0 } - } - - if (connectionStatus === 'connected') { - await mcpService.clearCache(workspaceId) + } else if ( + discoveryError !== null && + connectionStatus === 'disconnected' && + lastError === null && + !oauthAuthorizationRequired + ) { + lastError = discoveryError + toolCount = 0 } return createMcpSuccessResponse({ diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/mcp.tsx b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/mcp.tsx index c2769a36e54..92f3fc7f60a 100644 --- a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/mcp.tsx +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/mcp.tsx @@ -85,7 +85,12 @@ function ServerListItem({ onViewDetails, }: ServerListItemProps) { const transportLabel = formatTransportLabel(server.transport || 'http') - const toolsLabel = getServerToolsLabel(tools, server.connectionStatus, server.lastError) + const toolsLabel = getServerToolsLabel( + tools, + server.connectionStatus, + server.lastError, + server.authType + ) const hasConnectionIssue = server.connectionStatus === 'error' || server.connectionStatus === 'disconnected' @@ -380,6 +385,8 @@ export function MCP() { const refreshAction = getRefreshActionState({ mutationStatus: isCurrentRefresh ? refreshServerMutation.status : 'idle', connectionStatus: isCurrentRefresh ? refreshServerMutation.data?.status : undefined, + authType: server.authType, + error: isCurrentRefresh ? refreshServerMutation.data?.error : undefined, workflowsUpdated: isCurrentRefresh ? refreshServerMutation.data?.workflowsUpdated : undefined, }) @@ -427,7 +434,12 @@ export function MCP() {
Status

- {getServerToolsLabel([], server.connectionStatus, server.lastError)} + {getServerToolsLabel( + [], + server.connectionStatus, + server.lastError, + server.authType + )}

)} diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.test.ts b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.test.ts index a9eeb89513e..2ac09d8b300 100644 --- a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.test.ts @@ -31,6 +31,37 @@ describe('getRefreshActionState', () => { }) }) + it('shows Authorization required when an OAuth refresh finishes disconnected', () => { + expect( + getRefreshActionState({ + mutationStatus: 'success', + connectionStatus: 'disconnected', + authType: 'oauth', + workflowsUpdated: 0, + }) + ).toEqual({ + text: 'Authorization required', + textTone: 'error', + disabled: false, + }) + }) + + it('keeps Failed when a disconnected OAuth refresh has a concrete error', () => { + expect( + getRefreshActionState({ + mutationStatus: 'success', + connectionStatus: 'disconnected', + authType: 'oauth', + error: 'The MCP server took too long to respond and timed out', + workflowsUpdated: 0, + }) + ).toEqual({ + text: 'Failed', + textTone: 'error', + disabled: false, + }) + }) + it('preserves successful refresh feedback', () => { expect( getRefreshActionState({ diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.ts b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.ts index 28a6869921e..7f37d47d9d4 100644 --- a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.ts +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/refresh-action-state.ts @@ -1,10 +1,12 @@ import type { MutationStatus } from '@tanstack/react-query' -import type { RefreshMcpServerResult } from '@/lib/api/contracts/mcp' +import type { McpServer, RefreshMcpServerResult } from '@/lib/api/contracts/mcp' import type { SettingsAction } from '@/app/workspace/[workspaceId]/settings/components/settings-header/settings-header' interface RefreshActionStateInput { mutationStatus: MutationStatus connectionStatus?: RefreshMcpServerResult['status'] + authType?: McpServer['authType'] + error?: RefreshMcpServerResult['error'] workflowsUpdated?: number } @@ -13,12 +15,23 @@ type RefreshActionState = Pick export function getRefreshActionState({ mutationStatus, connectionStatus, + authType, + error, workflowsUpdated, }: RefreshActionStateInput): RefreshActionState { if (mutationStatus === 'pending') { return { text: 'Refreshing...', textTone: undefined, disabled: true } } + if ( + mutationStatus === 'success' && + connectionStatus === 'disconnected' && + authType === 'oauth' && + !error?.trim() + ) { + return { text: 'Authorization required', textTone: 'error', disabled: false } + } + if ( mutationStatus === 'error' || (mutationStatus === 'success' && connectionStatus !== 'connected') diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.test.ts b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.test.ts index e65dadf68e4..373ae480af4 100644 --- a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.test.ts @@ -12,12 +12,20 @@ describe('getServerToolsLabel', () => { expect(getServerToolsLabel([], 'error', null)).toBe('Unable to connect') }) - it('shows a disconnected state when OAuth was not completed', () => { - expect(getServerToolsLabel([], 'disconnected', null)).toBe('Not Connected') + it('shows that OAuth authorization is required when OAuth was not completed', () => { + expect(getServerToolsLabel([], 'disconnected', null, 'oauth')).toBe( + 'OAuth authorization required' + ) + }) + + it('keeps the generic disconnected state for non-OAuth servers', () => { + expect(getServerToolsLabel([], 'disconnected', null, 'headers')).toBe('Not Connected') }) it('shows the persisted error for disconnected connections', () => { - expect(getServerToolsLabel([], 'disconnected', 'Request timed out')).toBe('Request timed out') + expect(getServerToolsLabel([], 'disconnected', 'Request timed out', 'oauth')).toBe( + 'Request timed out' + ) }) it('continues showing discovered tools for healthy connections', () => { diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.ts b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.ts index f9d8d4896af..495cea5f9ff 100644 --- a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.ts +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/server-tools-label.ts @@ -7,14 +7,17 @@ interface NamedTool { export function getServerToolsLabel( tools: NamedTool[], connectionStatus?: McpServer['connectionStatus'], - lastError?: McpServer['lastError'] + lastError?: McpServer['lastError'], + authType?: McpServer['authType'] ): string { if (connectionStatus === 'error') { return lastError?.trim() || 'Unable to connect' } if (connectionStatus === 'disconnected') { - return lastError?.trim() || 'Not Connected' + return ( + lastError?.trim() || (authType === 'oauth' ? 'OAuth authorization required' : 'Not Connected') + ) } const count = tools.length diff --git a/apps/sim/lib/core/security/input-validation.server.ts b/apps/sim/lib/core/security/input-validation.server.ts index 7f81039622c..e3074a2accd 100644 --- a/apps/sim/lib/core/security/input-validation.server.ts +++ b/apps/sim/lib/core/security/input-validation.server.ts @@ -433,9 +433,21 @@ export function createPinnedLookup(resolvedIP: string): LookupFunction { * * The `Agent` is captured for the lifetime of the returned function, so repeated * calls (e.g. a provider tool loop) reuse its keep-alive connections. + * + * `allowH2` opts the pinned Agent into HTTP/2 (ALPN-negotiated, h1.1 fallback). + * It defaults to `false` to leave existing consumers unchanged. Enabling it does + * not weaken pinning: the pinned `connect.lookup` forces every connection on the + * Agent to `resolvedIP` regardless of authority, so h2 connection coalescing can + * never reach an address other than the validated one. */ -export function createPinnedFetch(resolvedIP: string): typeof fetch { - const dispatcher = new Agent({ connect: { lookup: createPinnedLookup(resolvedIP) } }) +export function createPinnedFetch( + resolvedIP: string, + options?: { allowH2?: boolean } +): typeof fetch { + const dispatcher = new Agent({ + allowH2: options?.allowH2 ?? false, + connect: { lookup: createPinnedLookup(resolvedIP) }, + }) const pinned = async (input: RequestInfo | URL, init?: RequestInit): Promise => { // double-cast-allowed: DOM RequestInfo/URL and undici fetch input types differ but are structurally compatible at runtime (Node's global fetch IS undici) diff --git a/apps/sim/lib/core/security/pinned-fetch.server.test.ts b/apps/sim/lib/core/security/pinned-fetch.server.test.ts index e07d0d412dd..0ecfa22cb9e 100644 --- a/apps/sim/lib/core/security/pinned-fetch.server.test.ts +++ b/apps/sim/lib/core/security/pinned-fetch.server.test.ts @@ -55,6 +55,18 @@ describe('createPinnedFetch', () => { expect(resolved).toEqual({ address: '203.0.113.10', family: 4 }) }) + it('defaults allowH2 to false so existing consumers are unchanged', () => { + createPinnedFetch('203.0.113.10') + const opts = capturedAgentOptions[0] as { allowH2?: boolean } + expect(opts.allowH2).toBe(false) + }) + + it('opts the Agent into HTTP/2 when allowH2 is requested', () => { + createPinnedFetch('203.0.113.10', { allowH2: true }) + const opts = capturedAgentOptions[0] as { allowH2?: boolean } + expect(opts.allowH2).toBe(true) + }) + it('uses IPv6 family when the validated IP is IPv6', async () => { createPinnedFetch('2606:4700:4700::1111') const { connect } = capturedAgentOptions[0] as { connect: { lookup: PinnedLookup } } diff --git a/apps/sim/lib/mcp/client.test.ts b/apps/sim/lib/mcp/client.test.ts index c568f425d61..198a67e5c4c 100644 --- a/apps/sim/lib/mcp/client.test.ts +++ b/apps/sim/lib/mcp/client.test.ts @@ -1,6 +1,7 @@ /** * @vitest-environment node */ +import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js' import { beforeEach, describe, expect, it, vi } from 'vitest' const { mockLogger, mockSdkConnect, mockSdkListTools } = vi.hoisted(() => ({ @@ -265,6 +266,55 @@ describe('McpClient notification handler', () => { expect(JSON.stringify(mockLogger.error.mock.calls)).not.toContain(secret) }) + it('does not misclassify rejected static headers as an OAuth authorization flow', async () => { + mockSdkConnect.mockRejectedValueOnce(new UnauthorizedError('Static token rejected')) + const client = new McpClient({ + config: { + ...createConfig(), + authType: 'headers', + headers: { Authorization: 'Bearer rejected-static-token' }, + }, + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await expect(client.connect()).rejects.toBeInstanceOf(UnauthorizedError) + + expect(mockLogger.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to connect'), + expect.objectContaining({ outcome: 'unauthorized' }) + ) + expect(client.getStatus().lastError).toBe('Authentication failed') + }) + + it('logs tools/list failures without echoed credentials or session identifiers', async () => { + const secret = 'opaque-tools-list-credential' + mockSdkListTools.mockRejectedValueOnce(new Error(`Upstream rejected ${secret}`)) + const client = new McpClient({ + config: { + ...createConfig(), + authType: 'headers', + headers: { 'X-Custom-Credential': secret }, + }, + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await client.connect() + await expect(client.listTools()).rejects.toThrow(secret) + + expect(mockLogger.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to list tools'), + expect.objectContaining({ + phase: 'tools/list', + serverId: 'server-1', + sessionIdPresent: true, + error: expect.objectContaining({ name: 'Error' }), + }) + ) + const logged = JSON.stringify(mockLogger.error.mock.calls) + expect(logged).not.toContain(secret) + expect(logged).not.toContain('test-session') + }) + it('passes configured headers for OAuth transports as well as header auth transports', () => { const authProvider = {} as unknown as NonNullable new McpClient({ diff --git a/apps/sim/lib/mcp/client.ts b/apps/sim/lib/mcp/client.ts index 50a90355c14..bc3d32c6dde 100644 --- a/apps/sim/lib/mcp/client.ts +++ b/apps/sim/lib/mcp/client.ts @@ -9,11 +9,11 @@ import { ToolListChangedNotificationSchema, } from '@modelcontextprotocol/sdk/types.js' import { createLogger } from '@sim/logger' -import { describeError, getErrorMessage } from '@sim/utils/errors' +import { getErrorMessage } from '@sim/utils/errors' import { getMaxExecutionTimeout } from '@/lib/core/execution-limits' -import { createPinnedFetch } from '@/lib/core/security/input-validation.server' -import { sanitizeForLogging } from '@/lib/core/security/redaction' +import { getMcpSafeErrorDiagnostics } from '@/lib/mcp/error-diagnostics' import { McpOauthRedirectRequired } from '@/lib/mcp/oauth' +import { createPinnedMcpFetch } from '@/lib/mcp/pinned-fetch' import { type McpClientOptions, McpConnectionError, @@ -43,10 +43,16 @@ type ConnectionOutcome = | 'cancelled' | 'error' -function classifyConnectionOutcome(error: unknown): ConnectionOutcome { - if (error instanceof McpOauthRedirectRequired || error instanceof UnauthorizedError) { +function classifyConnectionOutcome( + error: unknown, + authType: McpServerConfig['authType'] +): ConnectionOutcome { + if (error instanceof McpOauthRedirectRequired) { return 'authorization_required' } + if (error instanceof UnauthorizedError) { + return authType === 'oauth' ? 'authorization_required' : 'unauthorized' + } const message = getErrorMessage(error, '').toLowerCase() if (message.includes('connection attempt cancelled')) return 'cancelled' if (message.includes('timeout') || message.includes('timed out')) return 'timeout' @@ -92,7 +98,7 @@ export class McpClient { this.transport = new StreamableHTTPClientTransport(new URL(this.config.url), { authProvider: useOauth ? this.authProvider : undefined, requestInit: { headers: this.config.headers }, - ...(resolvedIP ? { fetch: createPinnedFetch(resolvedIP) } : {}), + ...(resolvedIP ? { fetch: createPinnedMcpFetch(resolvedIP) } : {}), }) this.client = new Client( @@ -109,7 +115,9 @@ export class McpClient { this.client.onerror = (error) => { logger.warn(`MCP transport error for ${this.config.name}`, { serverId: this.config.id, - error: sanitizeForLogging(getErrorMessage(error, 'Unknown transport error'), 200), + phase: 'transport', + sessionIdPresent: Boolean(this.transport.sessionId), + error: getMcpSafeErrorDiagnostics(error), }) } } @@ -178,25 +186,21 @@ export class McpClient { } catch (error) { this.isConnected = false const errorMessage = getErrorMessage(error, 'Unknown error') - const describedError = describeError(error) - const outcome = classifyConnectionOutcome(error) + const outcome = classifyConnectionOutcome(error, this.config.authType) logger.error(`Failed to connect to MCP server ${this.config.name}`, { ...diagnostics, durationMs: Date.now() - startedAt, - error: { - name: sanitizeForLogging(describedError.name, 100), - code: describedError.code ? sanitizeForLogging(describedError.code, 100) : undefined, - errno: describedError.errno ? sanitizeForLogging(describedError.errno, 100) : undefined, - syscall: describedError.syscall - ? sanitizeForLogging(describedError.syscall, 100) - : undefined, - }, + error: getMcpSafeErrorDiagnostics(error), outcome, }) if (outcome === 'authorization_required') { this.connectionStatus.lastError = undefined throw error } + if (error instanceof UnauthorizedError) { + this.connectionStatus.lastError = 'Authentication failed' + throw error + } this.connectionStatus.lastError = errorMessage throw new McpConnectionError(errorMessage, this.config.name) } @@ -236,6 +240,7 @@ export class McpClient { MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOTAL_TIMEOUT_MS ) const maxTotalTimeoutMs = MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOTAL_TIMEOUT_MS + const startedAt = Date.now() try { const result: ListToolsResult = await this.client.listTools(undefined, { @@ -265,7 +270,15 @@ export class McpClient { serverName: this.config.name, })) } catch (error) { - logger.error(`Failed to list tools from server ${this.config.name}:`, error) + logger.error(`Failed to list tools from server ${this.config.name}`, { + serverId: this.config.id, + phase: 'tools/list', + durationMs: Date.now() - startedAt, + idleTimeoutMs, + maxTotalTimeoutMs, + sessionIdPresent: Boolean(this.transport.sessionId), + error: getMcpSafeErrorDiagnostics(error), + }) throw error } } diff --git a/apps/sim/lib/mcp/error-diagnostics.test.ts b/apps/sim/lib/mcp/error-diagnostics.test.ts new file mode 100644 index 00000000000..2badd603c05 --- /dev/null +++ b/apps/sim/lib/mcp/error-diagnostics.test.ts @@ -0,0 +1,39 @@ +/** + * @vitest-environment node + */ + +import { describe, expect, it } from 'vitest' +import { getMcpSafeErrorDiagnostics } from './error-diagnostics' + +describe('getMcpSafeErrorDiagnostics', () => { + it('redacts and truncates structural fields without emitting message text', () => { + const secret = 'sk-abcdefghijklmnopqrstuvwxyz1234567890' + const error = Object.assign(new Error(`message contains ${secret}`), { + name: `Bearer ${secret}`, + code: `code-${'c'.repeat(200)}`, + errno: `api_key: "${secret}"`, + syscall: `Bearer ${secret}`, + sessionId: secret, + }) + + const diagnostics = getMcpSafeErrorDiagnostics(error) + + expect(diagnostics).toEqual({ + name: expect.any(String), + code: expect.any(String), + errno: expect.any(String), + syscall: expect.any(String), + }) + expect(diagnostics).not.toHaveProperty('message') + expect(diagnostics).not.toHaveProperty('causeChain') + expect(diagnostics).not.toHaveProperty('sessionId') + for (const value of Object.values(diagnostics)) { + expect(value).not.toContain(secret) + expect(value?.length).toBeLessThanOrEqual(100) + } + expect(diagnostics.name).toContain('[REDACTED]') + expect(diagnostics.errno).toContain('[REDACTED]') + expect(diagnostics.syscall).toContain('[REDACTED]') + expect(diagnostics.code).toHaveLength(100) + }) +}) diff --git a/apps/sim/lib/mcp/error-diagnostics.ts b/apps/sim/lib/mcp/error-diagnostics.ts new file mode 100644 index 00000000000..ce3b6bac7b4 --- /dev/null +++ b/apps/sim/lib/mcp/error-diagnostics.ts @@ -0,0 +1,25 @@ +import { describeError } from '@sim/utils/errors' +import { sanitizeForLogging } from '@/lib/core/security/redaction' + +const MAX_DIAGNOSTIC_FIELD_LENGTH = 100 + +export interface McpSafeErrorDiagnostics { + name: string + code: string | undefined + errno: string | undefined + syscall: string | undefined +} + +/** Returns bounded structural error fields without messages, causes, or session identifiers. */ +export function getMcpSafeErrorDiagnostics(error: unknown): McpSafeErrorDiagnostics { + const described = describeError(error) + const sanitizeOptional = (value: string | undefined) => + value === undefined ? undefined : sanitizeForLogging(value, MAX_DIAGNOSTIC_FIELD_LENGTH) + + return { + name: sanitizeForLogging(described.name, MAX_DIAGNOSTIC_FIELD_LENGTH), + code: sanitizeOptional(described.code), + errno: sanitizeOptional(described.errno), + syscall: sanitizeOptional(described.syscall), + } +} diff --git a/apps/sim/lib/mcp/oauth/probe.test.ts b/apps/sim/lib/mcp/oauth/probe.test.ts index d691f1178c7..ff5d96525a7 100644 --- a/apps/sim/lib/mcp/oauth/probe.test.ts +++ b/apps/sim/lib/mcp/oauth/probe.test.ts @@ -15,10 +15,8 @@ const { mockCreatePinnedFetch, mockCreateSsrfGuardedMcpFetch, mockPinnedFetch, m } }) -vi.mock('@/lib/core/security/input-validation.server', () => ({ - createPinnedFetch: mockCreatePinnedFetch, -})) vi.mock('@/lib/mcp/pinned-fetch', () => ({ + createPinnedMcpFetch: mockCreatePinnedFetch, createSsrfGuardedMcpFetch: mockCreateSsrfGuardedMcpFetch, })) diff --git a/apps/sim/lib/mcp/oauth/probe.ts b/apps/sim/lib/mcp/oauth/probe.ts index 4343f27870a..431480ae0f6 100644 --- a/apps/sim/lib/mcp/oauth/probe.ts +++ b/apps/sim/lib/mcp/oauth/probe.ts @@ -1,9 +1,8 @@ import { extractWWWAuthenticateParams } from '@modelcontextprotocol/sdk/client/auth.js' import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js' import { createLogger } from '@sim/logger' -import { createPinnedFetch } from '@/lib/core/security/input-validation.server' import { isLoopbackHostname } from '@/lib/core/utils/urls' -import { createSsrfGuardedMcpFetch } from '@/lib/mcp/pinned-fetch' +import { createPinnedMcpFetch, createSsrfGuardedMcpFetch } from '@/lib/mcp/pinned-fetch' import type { McpAuthType } from '@/lib/mcp/types' const logger = createLogger('McpOauthProbe') @@ -34,7 +33,7 @@ export async function detectMcpAuthType( } const probeFetch: FetchLike = resolvedIP - ? createPinnedFetch(resolvedIP) + ? createPinnedMcpFetch(resolvedIP) : createSsrfGuardedMcpFetch() const controller = new AbortController() diff --git a/apps/sim/lib/mcp/orchestration/server-lifecycle.test.ts b/apps/sim/lib/mcp/orchestration/server-lifecycle.test.ts new file mode 100644 index 00000000000..7b7058fa5bb --- /dev/null +++ b/apps/sim/lib/mcp/orchestration/server-lifecycle.test.ts @@ -0,0 +1,93 @@ +/** + * @vitest-environment node + */ +import { + auditMock, + dbChainMock, + dbChainMockFns, + drizzleOrmMock, + encryptionMock, + loggerMock, + posthogServerMock, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockClearCache, mockOauthCredsChanged } = vi.hoisted(() => ({ + mockClearCache: vi.fn(), + mockOauthCredsChanged: vi.fn(), +})) + +vi.mock('@sim/audit', () => auditMock) +vi.mock('@sim/db', () => ({ + ...dbChainMock, + mcpServers: schemaMock.mcpServers, +})) +vi.mock('@sim/db/schema', () => ({ + mcpServerOauth: schemaMock.mcpServerOauth, +})) +vi.mock('@sim/logger', () => loggerMock) +vi.mock('@sim/utils/id', () => ({ generateId: vi.fn() })) +vi.mock('drizzle-orm', () => drizzleOrmMock) +vi.mock('@/lib/core/security/encryption', () => encryptionMock) +vi.mock('@/lib/mcp/domain-check', () => ({ + McpDnsResolutionError: class extends Error {}, + McpDomainNotAllowedError: class extends Error {}, + McpSsrfError: class extends Error {}, + validateMcpDomain: vi.fn(), + validateMcpServerSsrf: vi.fn(), +})) +vi.mock('@/lib/mcp/oauth', () => ({ + detectMcpAuthType: vi.fn(), + oauthCredsChanged: mockOauthCredsChanged, + revokeMcpOauthTokens: vi.fn(), +})) +vi.mock('@/lib/mcp/service', () => ({ + mcpService: { clearCache: mockClearCache }, +})) +vi.mock('@/lib/mcp/utils', () => ({ generateMcpServerId: vi.fn() })) +vi.mock('@/lib/posthog/server', () => posthogServerMock) + +import { performUpdateMcpServer } from '@/lib/mcp/orchestration/server-lifecycle' + +describe('MCP server lifecycle orchestration', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + mockOauthCredsChanged.mockResolvedValue(false) + }) + + it('clears the workspace cache when an OAuth client ID implicitly changes the auth type', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([ + { + url: 'https://example.com/mcp', + authType: 'headers', + oauthClientId: 'client-1', + oauthClientSecret: null, + }, + ]) + dbChainMockFns.returning.mockResolvedValueOnce([ + { + id: 'server-1', + workspaceId: 'workspace-1', + name: 'Example', + transport: 'streamable-http', + url: 'https://example.com/mcp', + authType: 'oauth', + }, + ]) + + const result = await performUpdateMcpServer({ + workspaceId: 'workspace-1', + userId: 'user-1', + serverId: 'server-1', + oauthClientId: 'client-1', + oauthClientIdProvided: true, + }) + + expect(result.success).toBe(true) + expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ authType: 'oauth' })) + expect(mockClearCache).toHaveBeenCalledWith('workspace-1') + }) +}) diff --git a/apps/sim/lib/mcp/orchestration/server-lifecycle.ts b/apps/sim/lib/mcp/orchestration/server-lifecycle.ts index d8d2d73115a..83c82fec755 100644 --- a/apps/sim/lib/mcp/orchestration/server-lifecycle.ts +++ b/apps/sim/lib/mcp/orchestration/server-lifecycle.ts @@ -351,6 +351,7 @@ export async function performUpdateMcpServer( }) const shouldClearOauth = urlChanged || credsChanged const resolvedAuthType = (updateData.authType ?? currentServer.authType) as McpAuthType + const authTypeChanged = resolvedAuthType !== currentServer.authType if (shouldClearOauth && resolvedAuthType === 'oauth') { updateData.connectionStatus = 'disconnected' updateData.lastConnected = null @@ -384,6 +385,8 @@ export async function performUpdateMcpServer( const shouldClearCache = urlChanged || credsChanged || + params.transport !== undefined || + authTypeChanged || params.enabled !== undefined || params.headers !== undefined || params.timeout !== undefined || diff --git a/apps/sim/lib/mcp/pinned-fetch.test.ts b/apps/sim/lib/mcp/pinned-fetch.test.ts index 64354ee708b..3dc7701d6a7 100644 --- a/apps/sim/lib/mcp/pinned-fetch.test.ts +++ b/apps/sim/lib/mcp/pinned-fetch.test.ts @@ -31,7 +31,7 @@ describe('createSsrfGuardedMcpFetch', () => { await fetchLike('https://attacker.example/revoke', { method: 'POST' }) expect(mockValidateMcpServerSsrf).toHaveBeenCalledWith('https://attacker.example/revoke') - expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10') + expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10', { allowH2: true }) expect(sentinelFetch).toHaveBeenCalledWith('https://attacker.example/revoke', { method: 'POST', }) @@ -54,7 +54,7 @@ describe('createSsrfGuardedMcpFetch', () => { await fetchLike(new URL('https://attacker.example/discover')) expect(mockValidateMcpServerSsrf).toHaveBeenCalledWith('https://attacker.example/discover') - expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10') + expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10', { allowH2: true }) }) it('falls back to global fetch when validation returns no IP', async () => { diff --git a/apps/sim/lib/mcp/pinned-fetch.ts b/apps/sim/lib/mcp/pinned-fetch.ts index 3184a0da7a9..4188d348cac 100644 --- a/apps/sim/lib/mcp/pinned-fetch.ts +++ b/apps/sim/lib/mcp/pinned-fetch.ts @@ -2,6 +2,18 @@ import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js' import { createPinnedFetch } from '@/lib/core/security/input-validation.server' import { validateMcpServerSsrf } from '@/lib/mcp/domain-check' +/** + * Pinned fetch for all MCP traffic. MCP servers are commonly deployed behind + * HTTP/2 fronts (CDNs, cloud LBs), and undici's Agent is h1.1-only unless opted + * into h2 via ALPN — so every MCP connection enables it. This is the single + * source of the "MCP implies h2" decision; the base `createPinnedFetch` keeps + * its h1.1 default for non-MCP consumers. Pinning is unaffected: the pinned + * lookup forces the socket to `resolvedIP` regardless of negotiated protocol. + */ +export function createPinnedMcpFetch(resolvedIP: string): typeof fetch { + return createPinnedFetch(resolvedIP, { allowH2: true }) +} + /** * Builds a `FetchLike` that validates every outbound request URL against the * MCP SSRF policy before issuing it, then pins the connection to the resolved @@ -25,7 +37,7 @@ export function createSsrfGuardedMcpFetch(): FetchLike { return (async (url, init) => { const target = typeof url === 'string' ? url : url.href const resolvedIP = await validateMcpServerSsrf(target) - const pinnedFetch: FetchLike = resolvedIP ? createPinnedFetch(resolvedIP) : globalThis.fetch + const pinnedFetch: FetchLike = resolvedIP ? createPinnedMcpFetch(resolvedIP) : globalThis.fetch return pinnedFetch(url, init) }) satisfies FetchLike } diff --git a/apps/sim/lib/mcp/service.test.ts b/apps/sim/lib/mcp/service.test.ts index 7349b39e71f..11ee3db7d83 100644 --- a/apps/sim/lib/mcp/service.test.ts +++ b/apps/sim/lib/mcp/service.test.ts @@ -19,6 +19,7 @@ const { mockCacheAdapter, mockUpdateSet, mockUpdateReturning, + cacheStore, } = vi.hoisted(() => { const mockListTools = vi.fn() const mockConnect = vi.fn() @@ -28,6 +29,8 @@ const { // local .env points at (unreachable in CI/sandbox → hangs). Honors TTL via // an expiry timestamp so negative-cache assertions behave like production. const cacheStore = new Map() + const cacheMutations = new Map() + let nextMutationId = 0 const mockCacheAdapter = { get: vi.fn(async (key: string) => { const entry = cacheStore.get(key) @@ -44,12 +47,42 @@ const { delete: vi.fn(async (key: string) => { cacheStore.delete(key) }), + beginMutation: vi.fn(async (scopeKey: string) => { + const mutationId = Math.max(nextMutationId + 1, Date.now()) + nextMutationId = mutationId + cacheMutations.set(scopeKey, mutationId) + return mutationId + }), + applyMutationIfCurrent: vi.fn( + async ( + scopeKey: string, + mutationId: number, + setEntry: { key: string; tools: unknown[]; ttlMs: number } | null, + deleteKeys: string[] + ) => { + if (cacheMutations.get(scopeKey) !== mutationId) return false + if (setEntry) { + cacheStore.set(setEntry.key, { + tools: setEntry.tools, + expiry: Date.now() + setEntry.ttlMs, + }) + } + for (const key of deleteKeys) cacheStore.delete(key) + return true + } + ), clear: vi.fn(async () => { + for (const scopeKey of cacheMutations.keys()) { + const mutationId = Math.max(nextMutationId + 1, Date.now()) + nextMutationId = mutationId + cacheMutations.set(scopeKey, mutationId) + } cacheStore.clear() }), dispose: () => {}, } return { + cacheStore, mockCacheAdapter, MockMcpClient: vi.fn().mockImplementation( class { @@ -132,7 +165,7 @@ vi.mock('@/lib/mcp/storage', () => ({ getMcpCacheType: () => 'memory', })) -import { mcpService } from '@/lib/mcp/service' +import { getTimestampMillisecondBounds, mcpService } from '@/lib/mcp/service' import { McpOauthAuthorizationRequiredError } from '@/lib/mcp/types' const mockLogger = vi.mocked(loggerMock.createLogger).mock.results.at(-1)?.value @@ -170,6 +203,23 @@ function tool(name: string, serverId: string) { } } +describe('getTimestampMillisecondBounds', () => { + it('includes PostgreSQL sub-millisecond precision but excludes the next millisecond', () => { + const { startInclusive, endExclusive } = getTimestampMillisecondBounds( + '2026-01-01T00:00:00.123Z' + ) + const startMicroseconds = startInclusive.getTime() * 1_000 + const endMicroseconds = endExclusive.getTime() * 1_000 + const isWithinBounds = (candidateMicroseconds: number) => + candidateMicroseconds >= startMicroseconds && candidateMicroseconds < endMicroseconds + + // PostgreSQL can retain any of these extra microseconds even though the + // JavaScript Date used as the generation token is truncated to .123. + expect(isWithinBounds(startMicroseconds + 999)).toBe(true) + expect(isWithinBounds(endMicroseconds)).toBe(false) + }) +}) + describe('McpService.discoverTools per-server caching', () => { beforeEach(async () => { vi.clearAllMocks() @@ -279,6 +329,15 @@ describe('McpService.discoverTools per-server caching', () => { await mcpService.clearCache(WORKSPACE_ID) + const [discoveryUpdate, invalidationUpdate] = mockUpdateSet.mock.calls.map(([update]) => update) + expect(invalidationUpdate).toEqual({ + toolCount: 0, + lastToolsRefresh: expect.any(Date), + }) + expect(invalidationUpdate.lastToolsRefresh.getTime()).toBeGreaterThan( + discoveryUpdate.lastToolsRefresh.getTime() + ) + mockListTools.mockClear() mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) await mcpService.discoverTools(USER_ID, WORKSPACE_ID) @@ -341,7 +400,7 @@ describe('McpService.discoverTools per-server caching', () => { statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, }), ]) - mockListTools.mockRejectedValueOnce( + mockConnect.mockRejectedValueOnce( new UnauthorizedError(`Rejected Authorization: ${reflectedCredential}`) ) @@ -356,15 +415,23 @@ describe('McpService.discoverTools per-server caching', () => { statusConfig: { consecutiveFailures: 1, lastSuccessfulDiscovery: null }, }) ) - expect(mockCacheAdapter.set).toHaveBeenCalledWith( - `workspace:${WORKSPACE_ID}:server:mcp-a:failure`, - [], - expect.any(Number) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledWith( + `workspace:${WORKSPACE_ID}:server:mcp-a`, + expect.any(Number), + { + key: `workspace:${WORKSPACE_ID}:server:mcp-a:failure`, + tools: [], + ttlMs: expect.any(Number), + }, + [`workspace:${WORKSPACE_ID}:server:mcp-a`] ) }) expect(JSON.stringify(mockUpdateSet.mock.calls)).not.toContain(reflectedCredential) - expect(JSON.stringify(mockCacheAdapter.set.mock.calls)).not.toContain(reflectedCredential) + expect(JSON.stringify(mockCacheAdapter.applyMutationIfCurrent.mock.calls)).not.toContain( + reflectedCredential + ) expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) + expect(mockListTools).not.toHaveBeenCalled() mockListTools.mockClear() const second = await mcpService.discoverTools(USER_ID, WORKSPACE_ID) @@ -387,10 +454,11 @@ describe('McpService.discoverTools per-server caching', () => { }) ) }) - expect(mockCacheAdapter.set).not.toHaveBeenCalledWith( - `workspace:${WORKSPACE_ID}:server:mcp-a:failure`, - [], - expect.any(Number) + expect(mockCacheAdapter.applyMutationIfCurrent).not.toHaveBeenCalledWith( + expect.anything(), + expect.anything(), + expect.objectContaining({ key: `workspace:${WORKSPACE_ID}:server:mcp-a:failure` }), + expect.anything() ) mockResolveEnvVars.mockClear() @@ -431,6 +499,351 @@ describe('McpService.discoverTools per-server caching', () => { expect(mockListTools).not.toHaveBeenCalled() }) + it('retries mutation ownership before publishing discovery state', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 }) + mockCacheAdapter.beginMutation.mockRejectedValueOnce(new Error('cache ordering unavailable')) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + const tools = await mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + + expect(tools).toEqual([tool('a1', 'mcp-a')]) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('a1', 'mcp-a')]) + expect(cacheStore.has(failureKey)).toBe(false) + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(2) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledWith( + serverKey, + expect.any(Number), + { key: serverKey, tools: [tool('a1', 'mcp-a')], ttlMs: expect.any(Number) }, + [failureKey] + ) + expect(mockCacheAdapter.set).not.toHaveBeenCalled() + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).toHaveBeenCalledWith( + expect.objectContaining({ connectionStatus: 'connected', toolCount: 1 }) + ) + }) + + it('acquires a fresh mutation on each discovery retry so a retried result still publishes', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockCacheAdapter.beginMutation.mockClear() + mockListTools + .mockRejectedValueOnce(new Error('Request timed out')) + .mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + const tools = await mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + + expect(tools).toEqual([tool('a1', 'mcp-a')]) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('a1', 'mcp-a')]) + // One begin per attempt: the retry publishes under a current ownership id + // instead of a stale pre-loop id that a concurrent clearCache could supersede. + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(2) + }) + + it('keeps an older ordered publisher from superseding a retry-acquired mutation', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let resolveOlder: ((tools: ReturnType[]) => void) | undefined + mockListTools + .mockReturnValueOnce( + new Promise((resolve) => { + resolveOlder = resolve + }) + ) + .mockResolvedValueOnce([tool('new-tool', 'mcp-a')]) + + const older = mcpService.discoverServerToolsWithMetadata(USER_ID, 'mcp-a', WORKSPACE_ID, false) + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1)) + + mockCacheAdapter.beginMutation.mockRejectedValueOnce(new Error('transient ordering failure')) + const newer = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await expect(newer).resolves.toEqual([tool('new-tool', 'mcp-a')]) + + resolveOlder?.([tool('old-tool', 'mcp-a')]) + await expect(older).resolves.toEqual({ + tools: [tool('new-tool', 'mcp-a')], + state: 'winner-cache', + }) + + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('new-tool', 'mcp-a')]) + expect(mockUpdateSet).toHaveBeenCalledTimes(1) + }) + + it('uses the cache mutation token to order database publication after a begin retry', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + try { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let rejectOlderBegin: ((error: Error) => void) | undefined + mockCacheAdapter.beginMutation.mockImplementationOnce( + () => + new Promise((_resolve, reject) => { + rejectOlderBegin = reject + }) + ) + mockListTools + .mockRejectedValueOnce(new Error('Later-started discovery failed')) + .mockResolvedValueOnce([tool('retry-winner', 'mcp-a')]) + + vi.setSystemTime(new Date('2030-02-01T00:00:00.000Z')) + const older = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false) + await vi.waitFor(() => expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(1)) + + const laterMutationTime = new Date('2030-02-01T00:00:01.000Z') + vi.setSystemTime(laterMutationTime) + const later = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await expect(later).rejects.toThrow('Later-started discovery failed') + + const retriedMutationTime = new Date('2030-02-01T00:00:02.000Z') + vi.setSystemTime(retriedMutationTime) + rejectOlderBegin?.(new Error('Transient mutation start failure')) + await expect(older).resolves.toEqual([tool('retry-winner', 'mcp-a')]) + + const publications = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.lastToolsRefresh) + expect(publications.map((update) => update.connectionStatus)).toEqual([ + 'disconnected', + 'connected', + ]) + expect(publications.map((update) => update.lastToolsRefresh)).toEqual([ + laterMutationTime, + retriedMutationTime, + ]) + expect(cacheStore.get(`workspace:${WORKSPACE_ID}:server:mcp-a`)?.tools).toEqual([ + tool('retry-winner', 'mcp-a'), + ]) + expect(cacheStore.has(`workspace:${WORKSPACE_ID}:server:mcp-a:failure`)).toBe(false) + } finally { + vi.useRealTimers() + } + }) + + it('returns live tools without publishing when cache ownership is unavailable', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let resolveOlder: ((tools: ReturnType[]) => void) | undefined + mockListTools + .mockReturnValueOnce( + new Promise((resolve) => { + resolveOlder = resolve + }) + ) + .mockResolvedValueOnce([tool('unowned-new-tool', 'mcp-a')]) + + const older = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false) + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1)) + + mockCacheAdapter.beginMutation + .mockRejectedValueOnce(new Error('ordering unavailable')) + .mockRejectedValueOnce(new Error('ordering unavailable')) + const newer = mcpService.discoverServerToolsWithMetadata(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await expect(newer).resolves.toEqual({ + tools: [tool('unowned-new-tool', 'mcp-a')], + state: 'unavailable', + }) + + resolveOlder?.([tool('owned-old-tool', 'mcp-a')]) + await expect(older).resolves.toEqual([tool('owned-old-tool', 'mcp-a')]) + + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('owned-old-tool', 'mcp-a')]) + expect(mockCacheAdapter.set).not.toHaveBeenCalled() + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).toHaveBeenCalledTimes(1) + }) + + it('returns live tools but skips publication when mutation ownership stays unavailable', async () => { + const reflectedCredential = 'opaque-cache-provider-message' + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockCacheAdapter.beginMutation + .mockRejectedValueOnce(new Error(reflectedCredential)) + .mockRejectedValueOnce(new Error(reflectedCredential)) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + await expect( + mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).resolves.toEqual([tool('a1', 'mcp-a')]) + + expect(mockCacheAdapter.applyMutationIfCurrent).not.toHaveBeenCalled() + expect(mockCacheAdapter.set).not.toHaveBeenCalled() + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).not.toHaveBeenCalled() + expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) + }) + + it('returns bulk live tools without publication when mutation ownership is unavailable', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockCacheAdapter.beginMutation + .mockRejectedValueOnce(new Error('cache ordering unavailable')) + .mockRejectedValueOnce(new Error('cache ordering unavailable')) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + await expect(mcpService.discoverTools(USER_ID, WORKSPACE_ID, true)).resolves.toEqual([ + tool('a1', 'mcp-a'), + ]) + + expect(mockCacheAdapter.applyMutationIfCurrent).not.toHaveBeenCalled() + expect(mockUpdateSet).not.toHaveBeenCalled() + }) + + it('returns live tools but skips publication when the atomic cache transition fails', async () => { + const reflectedCredential = 'opaque-atomic-cache-provider-message' + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockCacheAdapter.applyMutationIfCurrent.mockRejectedValueOnce(new Error(reflectedCredential)) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + await expect( + mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).resolves.toEqual([tool('a1', 'mcp-a')]) + + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(1) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(1) + expect(mockUpdateSet).not.toHaveBeenCalled() + expect(mockLogger?.warn).toHaveBeenCalledWith( + 'Failed to atomically update cache for server mcp-a', + expect.objectContaining({ + workspaceId: WORKSPACE_ID, + error: expect.objectContaining({ name: 'Error' }), + }) + ) + expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) + }) + + it('preserves cache and database state when invalidation cannot be ordered', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + cacheStore.set(serverKey, { + tools: [tool('stale-tool', 'mcp-a')], + expiry: Date.now() + 60_000, + }) + cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 }) + for (let attempt = 0; attempt < 6; attempt++) { + mockCacheAdapter.beginMutation.mockRejectedValueOnce(new Error('cache ordering unavailable')) + } + + await mcpService.clearCache(WORKSPACE_ID) + + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('stale-tool', 'mcp-a')]) + expect(cacheStore.has(failureKey)).toBe(true) + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(6) + expect(mockCacheAdapter.applyMutationIfCurrent).not.toHaveBeenCalled() + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).not.toHaveBeenCalled() + }) + + it('reacquires ownership after a transient atomic invalidation failure', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + cacheStore.set(serverKey, { + tools: [tool('stale-tool', 'mcp-a')], + expiry: Date.now() + 60_000, + }) + cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 }) + mockCacheAdapter.applyMutationIfCurrent.mockRejectedValueOnce( + new Error('atomic invalidation unavailable') + ) + + await mcpService.clearCache(WORKSPACE_ID) + + const firstMutationId = mockCacheAdapter.applyMutationIfCurrent.mock.calls[0][1] + const successfulMutationId = mockCacheAdapter.applyMutationIfCurrent.mock.calls[1][1] + expect(successfulMutationId).toBeGreaterThan(firstMutationId) + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(2) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(2) + expect(cacheStore.has(serverKey)).toBe(false) + expect(cacheStore.has(failureKey)).toBe(false) + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).toHaveBeenCalledWith({ + toolCount: 0, + lastToolsRefresh: new Date(successfulMutationId), + }) + }) + + it('preserves a newer winner acquired between invalidation retry begin and apply', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + const winner = tool('winner-tool', 'mcp-a') + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + cacheStore.set(serverKey, { + tools: [tool('stale-tool', 'mcp-a')], + expiry: Date.now() + 60_000, + }) + cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 }) + const defaultApply = mockCacheAdapter.applyMutationIfCurrent.getMockImplementation() + mockCacheAdapter.applyMutationIfCurrent + .mockRejectedValueOnce(new Error('atomic invalidation unavailable')) + .mockImplementationOnce(async (scopeKey, mutationId, setEntry, deleteKeys) => { + const newerMutationId = await mockCacheAdapter.beginMutation(scopeKey) + await defaultApply?.( + scopeKey, + newerMutationId, + { key: serverKey, tools: [winner], ttlMs: 60_000 }, + [failureKey] + ) + return defaultApply?.(scopeKey, mutationId, setEntry, deleteKeys) ?? false + }) + + await mcpService.clearCache(WORKSPACE_ID) + + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(2) + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3) + expect(cacheStore.get(serverKey)?.tools).toEqual([winner]) + expect(cacheStore.has(failureKey)).toBe(false) + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).not.toHaveBeenCalled() + }) + + it('bounds persistent atomic invalidation failures without raw deletes', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + const reflectedCredential = 'opaque-cache-provider-message' + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + cacheStore.set(serverKey, { + tools: [tool('stale-tool', 'mcp-a')], + expiry: Date.now() + 60_000, + }) + cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 }) + for (let attempt = 0; attempt < 3; attempt++) { + mockCacheAdapter.applyMutationIfCurrent.mockRejectedValueOnce(new Error(reflectedCredential)) + } + + await mcpService.clearCache(WORKSPACE_ID) + + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(3) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('stale-tool', 'mcp-a')]) + expect(cacheStore.has(failureKey)).toBe(true) + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).not.toHaveBeenCalled() + expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) + }) + + it('preserves a newer cache winner when invalidation is superseded', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + const winner = tool('winner-tool', 'mcp-a') + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + cacheStore.set(serverKey, { tools: [winner], expiry: Date.now() + 60_000 }) + cacheStore.set(failureKey, { tools: [], expiry: Date.now() + 60_000 }) + mockCacheAdapter.applyMutationIfCurrent.mockResolvedValueOnce(false) + + await mcpService.clearCache(WORKSPACE_ID) + + expect(cacheStore.get(serverKey)?.tools).toEqual([winner]) + expect(cacheStore.has(failureKey)).toBe(true) + expect(mockCacheAdapter.delete).not.toHaveBeenCalled() + expect(mockUpdateSet).not.toHaveBeenCalled() + }) + it('does not negative-cache OAuth-required errors', async () => { mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) mockListTools.mockRejectedValueOnce(new McpOauthAuthorizationRequiredError('mcp-a', 'A')) @@ -472,6 +885,623 @@ describe('McpService.discoverTools per-server caching', () => { ) }) + it('persists an allowlisted message when an upstream error reflects a custom credential', async () => { + const reflectedCredential = 'opaque-custom-header-value' + mockGetWorkspaceServersRows.mockResolvedValue([ + dbRow('mcp-a', 'A', { + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }), + ]) + mockListTools.mockRejectedValueOnce( + new Error(`Provider rejected X-Custom-Credential: ${reflectedCredential}`) + ) + + await expect(mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID)).rejects.toThrow( + reflectedCredential + ) + + expect(mockUpdateSet).toHaveBeenCalledWith( + expect.objectContaining({ + connectionStatus: 'disconnected', + lastError: 'Connection failed', + toolCount: 0, + }) + ) + expect(JSON.stringify(mockUpdateSet.mock.calls)).not.toContain(reflectedCredential) + }) + + it('does not return or cache tools discovered from a stale server configuration', async () => { + mockGetWorkspaceServersRows.mockResolvedValueOnce([dbRow('mcp-a', 'A')]).mockResolvedValueOnce([ + dbRow('mcp-a', 'A', { + url: 'https://changed-config.example.com/mcp', + updatedAt: new Date('2026-01-01T00:00:01Z'), + }), + ]) + mockListTools.mockResolvedValueOnce([tool('stale-tool', 'mcp-a')]) + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + mockUpdateSet.mockReturnValueOnce({ + where: vi.fn().mockReturnValue({ returning: vi.fn().mockResolvedValue([]) }), + }) + + const tools = await mcpService.discoverTools(USER_ID, WORKSPACE_ID, true) + + expect(tools).toEqual([]) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledWith( + expect.anything(), + expect.anything(), + { + key: `workspace:${WORKSPACE_ID}:server:mcp-a`, + tools: [tool('stale-tool', 'mcp-a')], + ttlMs: expect.any(Number), + }, + [`workspace:${WORKSPACE_ID}:server:mcp-a:failure`] + ) + expect(cacheStore.has(serverKey)).toBe(false) + }) + + it('publishes successful discovery after repeated metadata-only renames win the database CAS', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const original = dbRow('mcp-a', 'A') + const renamed = dbRow('mcp-a', 'Renamed A', { + description: 'Updated display-only description', + updatedAt: new Date('2026-01-01T00:00:01Z'), + }) + const renamedAgain = dbRow('mcp-a', 'Renamed A Again', { + description: 'Another display-only description', + updatedAt: new Date('2026-01-01T00:00:02Z'), + }) + mockGetWorkspaceServersRows + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([renamed]) + .mockResolvedValueOnce([renamedAgain]) + .mockResolvedValue([renamedAgain]) + mockListTools.mockResolvedValueOnce([tool('still-valid', 'mcp-a')]) + mockUpdateReturning + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ id: 'mcp-a' }]) + + await expect( + mcpService.discoverServerToolsWithMetadata(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).resolves.toEqual({ + tools: [tool('still-valid', 'mcp-a')], + state: 'published', + publicationOrder: expect.any(Date), + }) + + const successfulStatusWrites = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.connectionStatus === 'connected') + expect(successfulStatusWrites).toHaveLength(3) + expect(successfulStatusWrites.at(-1)).toEqual( + expect.objectContaining({ connectionStatus: 'connected', toolCount: 1 }) + ) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('still-valid', 'mcp-a')]) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenLastCalledWith( + serverKey, + expect.any(Number), + null, + [] + ) + }) + + it('keeps valid live tools when successful status publication retries are exhausted', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockListTools.mockResolvedValueOnce([tool('still-valid', 'mcp-a')]) + mockUpdateReturning.mockResolvedValue([]) + + await expect( + mcpService.discoverServerToolsWithMetadata(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).resolves.toEqual({ + tools: [tool('still-valid', 'mcp-a')], + state: 'unavailable', + }) + + const successfulStatusWrites = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.connectionStatus === 'connected') + expect(successfulStatusWrites).toHaveLength(4) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('still-valid', 'mcp-a')]) + }) + + it('publishes a failed discovery after repeated metadata-only renames win the database CAS', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + const original = dbRow('mcp-a', 'A', { + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }) + const renamed = dbRow('mcp-a', 'Renamed A', { + description: 'Updated display-only description', + updatedAt: new Date('2026-01-01T00:00:01Z'), + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }) + const renamedAgain = dbRow('mcp-a', 'Renamed A Again', { + description: 'Another display-only description', + updatedAt: new Date('2026-01-01T00:00:02Z'), + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }) + mockGetWorkspaceServersRows + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([renamed]) + .mockResolvedValueOnce([renamed]) + .mockResolvedValueOnce([renamed]) + .mockResolvedValueOnce([renamed]) + .mockResolvedValueOnce([renamedAgain]) + .mockResolvedValueOnce([renamedAgain]) + .mockResolvedValueOnce([renamedAgain]) + .mockResolvedValueOnce([renamedAgain]) + .mockResolvedValue([renamedAgain]) + mockUpdateReturning + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ id: 'mcp-a' }]) + mockListTools.mockRejectedValueOnce(new Error('Permanent discovery failure')) + cacheStore.set(serverKey, { + tools: [tool('previous-tool', 'mcp-a')], + expiry: Date.now() + 60_000, + }) + + await expect( + mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).rejects.toThrow('Permanent discovery failure') + + const failureStatusWrites = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.lastError === 'Connection failed') + expect(failureStatusWrites).toHaveLength(7) + expect(failureStatusWrites.at(-1)).toEqual( + expect.objectContaining({ connectionStatus: 'disconnected', toolCount: 0 }) + ) + expect(cacheStore.has(serverKey)).toBe(false) + expect(cacheStore.has(failureKey)).toBe(true) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledWith( + serverKey, + expect.any(Number), + null, + [] + ) + }) + + it('publishes OAuth pending after repeated metadata-only renames win the database CAS', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const original = dbRow('mcp-a', 'A') + const renamed = dbRow('mcp-a', 'Renamed A', { + description: 'Updated display-only description', + updatedAt: new Date('2026-01-01T00:00:01Z'), + }) + const renamedAgain = dbRow('mcp-a', 'Renamed A Again', { + description: 'Another display-only description', + updatedAt: new Date('2026-01-01T00:00:02Z'), + }) + mockGetWorkspaceServersRows + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([renamed]) + .mockResolvedValueOnce([renamedAgain]) + .mockResolvedValue([renamedAgain]) + mockUpdateReturning + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ id: 'mcp-a' }]) + mockListTools.mockRejectedValueOnce(new McpOauthAuthorizationRequiredError('mcp-a', 'A')) + cacheStore.set(serverKey, { + tools: [tool('previous-tool', 'mcp-a')], + expiry: Date.now() + 60_000, + }) + + await expect( + mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).rejects.toThrow('OAuth authorization required') + + const oauthStatusWrites = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.connectionStatus === 'disconnected' && update.lastError === null) + expect(oauthStatusWrites).toHaveLength(3) + expect(cacheStore.has(serverKey)).toBe(false) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledWith( + serverKey, + expect.any(Number), + null, + [] + ) + }) + + it('does not retry failed status publication after a connection-config edit', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + const original = dbRow('mcp-a', 'A', { + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }) + const reconfigured = dbRow('mcp-a', 'A', { + url: 'https://changed-config.example.com/mcp', + updatedAt: new Date('2026-01-01T00:00:01Z'), + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }) + mockGetWorkspaceServersRows + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([original]) + .mockResolvedValue([reconfigured]) + mockUpdateReturning.mockResolvedValue([]) + mockListTools.mockRejectedValueOnce(new Error('Permanent discovery failure')) + + await expect( + mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).rejects.toThrow('Permanent discovery failure') + + const failureStatusWrites = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.lastError === 'Connection failed') + expect(failureStatusWrites).toHaveLength(3) + expect(cacheStore.has(failureKey)).toBe(false) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenLastCalledWith( + serverKey, + expect.any(Number), + null, + [failureKey] + ) + }) + + it('does not retry OAuth status publication after a connection-config edit', async () => { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const original = dbRow('mcp-a', 'A') + const reconfigured = dbRow('mcp-a', 'A', { + url: 'https://changed-config.example.com/mcp', + updatedAt: new Date('2026-01-01T00:00:01Z'), + }) + mockGetWorkspaceServersRows + .mockResolvedValueOnce([original]) + .mockResolvedValueOnce([reconfigured]) + .mockResolvedValue([reconfigured]) + mockUpdateReturning.mockResolvedValue([]) + mockListTools.mockRejectedValueOnce(new McpOauthAuthorizationRequiredError('mcp-a', 'A')) + + await expect( + mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + ).rejects.toThrow('OAuth authorization required') + + const oauthStatusWrites = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.connectionStatus === 'disconnected' && update.lastError === null) + expect(oauthStatusWrites).toHaveLength(1) + expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenLastCalledWith( + serverKey, + expect.any(Number), + null, + [] + ) + }) + + it('supersedes an older discovery before it can publish status', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + try { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let resolveOlder: ((tools: ReturnType[]) => void) | undefined + let resolveNewer: ((tools: ReturnType[]) => void) | undefined + mockListTools + .mockReturnValueOnce( + new Promise((resolve) => { + resolveOlder = resolve + }) + ) + .mockReturnValueOnce( + new Promise((resolve) => { + resolveNewer = resolve + }) + ) + + const olderStartedAt = new Date('2030-02-01T00:00:00.000Z') + vi.setSystemTime(olderStartedAt) + const older = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false) + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1)) + + const newerStartedAt = new Date('2030-02-01T00:00:01.000Z') + vi.setSystemTime(newerStartedAt) + const newer = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(2)) + + vi.setSystemTime(new Date('2030-02-01T00:00:02.000Z')) + resolveOlder?.([tool('old-tool', 'mcp-a')]) + await expect(older).resolves.toEqual([]) + + vi.setSystemTime(new Date('2030-02-01T00:00:03.000Z')) + resolveNewer?.([tool('new-tool', 'mcp-a')]) + await expect(newer).resolves.toEqual([tool('new-tool', 'mcp-a')]) + + const publishedRefreshTimes = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.connectionStatus === 'connected') + .map((update) => update.lastToolsRefresh) + expect(publishedRefreshTimes).toHaveLength(1) + expect(publishedRefreshTimes[0].getTime()).toBeGreaterThanOrEqual(newerStartedAt.getTime()) + } finally { + vi.useRealTimers() + } + }) + + it('does not write status after its cache mutation is superseded', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let resolveList: ((tools: ReturnType[]) => void) | undefined + mockListTools.mockReturnValueOnce( + new Promise((resolve) => { + resolveList = resolve + }) + ) + + let releaseOlderCacheMutation: (() => void) | undefined + const olderCacheMutationGate = new Promise((resolve) => { + releaseOlderCacheMutation = resolve + }) + const defaultApply = mockCacheAdapter.applyMutationIfCurrent.getMockImplementation() + mockCacheAdapter.applyMutationIfCurrent.mockImplementationOnce( + async ( + scopeKey: string, + mutationId: number, + setEntry: { key: string; tools: unknown[]; ttlMs: number } | null, + deleteKeys: string[] + ) => { + await olderCacheMutationGate + return defaultApply?.(scopeKey, mutationId, setEntry, deleteKeys) ?? false + } + ) + + const discovery = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false) + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1)) + resolveList?.([tool('superseded-tool', 'mcp-a')]) + await vi.waitFor(() => expect(mockCacheAdapter.applyMutationIfCurrent).toHaveBeenCalledTimes(1)) + + await mockCacheAdapter.beginMutation(`workspace:${WORKSPACE_ID}:server:mcp-a`) + releaseOlderCacheMutation?.() + + await expect(discovery).resolves.toEqual([]) + expect(mockUpdateSet).not.toHaveBeenCalled() + expect(cacheStore.has(`workspace:${WORKSPACE_ID}:server:mcp-a`)).toBe(false) + }) + + it('waits for bulk discovery status publication before returning', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) + + let releaseStatus: (() => void) | undefined + const pendingStatus = new Promise((resolve) => { + releaseStatus = resolve + }) + mockUpdateSet.mockReturnValueOnce({ + where: vi.fn().mockReturnValue({ + returning: vi.fn().mockReturnValue(pendingStatus.then(() => [{ id: 'mcp-a' }])), + }), + }) + + let settled = false + const discovery = mcpService.discoverTools(USER_ID, WORKSPACE_ID, true).finally(() => { + settled = true + }) + + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1)) + await Promise.resolve() + expect(settled).toBe(false) + + releaseStatus?.() + await expect(discovery).resolves.toEqual([tool('a1', 'mcp-a')]) + }) + + it('publishes a newer invalidation barrier while successful status publication is pending', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + try { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const discoveryStartedAt = new Date('2100-02-01T00:00:00.000Z') + const invalidatedAt = new Date('2100-02-01T00:00:01.000Z') + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockListTools.mockResolvedValueOnce([tool('stale-tool', 'mcp-a')]) + + let releaseDiscoveryStatus: ((rows: { id: string }[]) => void) | undefined + mockUpdateReturning.mockReturnValueOnce( + new Promise((resolve) => { + releaseDiscoveryStatus = resolve + }) + ) + + vi.setSystemTime(discoveryStartedAt) + const discovery = mcpService.discoverServerToolsWithMetadata( + USER_ID, + 'mcp-a', + WORKSPACE_ID, + true + ) + await vi.waitFor(() => + expect(mockUpdateSet).toHaveBeenCalledWith( + expect.objectContaining({ + connectionStatus: 'connected', + lastToolsRefresh: discoveryStartedAt, + }) + ) + ) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('stale-tool', 'mcp-a')]) + + vi.setSystemTime(invalidatedAt) + await mcpService.clearCache(WORKSPACE_ID) + + const invalidationUpdate = mockUpdateSet.mock.calls + .map(([update]) => update) + .find( + (update) => + update.toolCount === 0 && update.lastToolsRefresh?.getTime() === invalidatedAt.getTime() + ) + expect(invalidationUpdate).toEqual({ + toolCount: 0, + lastToolsRefresh: invalidatedAt, + }) + expect(cacheStore.has(serverKey)).toBe(false) + + // The real lastToolsRefresh predicate rejects this older publication + // after the invalidation barrier wins the database race. + releaseDiscoveryStatus?.([]) + await expect(discovery).resolves.toEqual({ tools: [], state: 'superseded' }) + } finally { + vi.useRealTimers() + } + }) + + it('publishes a newer invalidation barrier while failed status publication is pending', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + try { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const failureKey = `${serverKey}:failure` + const discoveryStartedAt = new Date('2100-03-01T00:00:00.000Z') + const invalidatedAt = new Date('2100-03-01T00:00:01.000Z') + mockGetWorkspaceServersRows.mockResolvedValue([ + dbRow('mcp-a', 'A', { + statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null }, + }), + ]) + mockListTools.mockRejectedValueOnce(new Error('Permanent discovery failure')) + + let releaseDiscoveryStatus: ((rows: { id: string }[]) => void) | undefined + mockUpdateReturning + .mockReturnValueOnce( + new Promise((resolve) => { + releaseDiscoveryStatus = resolve + }) + ) + .mockResolvedValueOnce([{ id: 'mcp-a' }]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]) + + vi.setSystemTime(discoveryStartedAt) + const discovery = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await vi.waitFor(() => + expect(mockUpdateSet).toHaveBeenCalledWith( + expect.objectContaining({ + connectionStatus: 'disconnected', + lastError: 'Connection failed', + lastToolsRefresh: discoveryStartedAt, + }) + ) + ) + expect(cacheStore.has(failureKey)).toBe(true) + + vi.setSystemTime(invalidatedAt) + await mcpService.clearCache(WORKSPACE_ID) + + expect(mockUpdateSet).toHaveBeenCalledWith({ + toolCount: 0, + lastToolsRefresh: invalidatedAt, + }) + expect(cacheStore.has(serverKey)).toBe(false) + expect(cacheStore.has(failureKey)).toBe(false) + + releaseDiscoveryStatus?.([]) + await expect(discovery).rejects.toThrow('Permanent discovery failure') + expect( + mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.lastError === 'Connection failed') + .every((update) => update.lastToolsRefresh?.getTime() === discoveryStartedAt.getTime()) + ).toBe(true) + } finally { + vi.useRealTimers() + } + }) + + it('publishes a newer invalidation barrier while OAuth status publication is pending', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + try { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + const discoveryStartedAt = new Date('2100-04-01T00:00:00.000Z') + const invalidatedAt = new Date('2100-04-01T00:00:01.000Z') + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + mockListTools.mockRejectedValueOnce(new McpOauthAuthorizationRequiredError('mcp-a', 'A')) + + let releaseDiscoveryStatus: ((rows: { id: string }[]) => void) | undefined + mockUpdateReturning.mockReturnValueOnce( + new Promise((resolve) => { + releaseDiscoveryStatus = resolve + }) + ) + + vi.setSystemTime(discoveryStartedAt) + const discovery = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await vi.waitFor(() => + expect(mockUpdateSet).toHaveBeenCalledWith( + expect.objectContaining({ + connectionStatus: 'disconnected', + lastError: null, + lastToolsRefresh: discoveryStartedAt, + }) + ) + ) + + vi.setSystemTime(invalidatedAt) + await mcpService.clearCache(WORKSPACE_ID) + + const invalidationUpdate = mockUpdateSet.mock.calls + .map(([update]) => update) + .find( + (update) => + update.toolCount === 0 && update.lastToolsRefresh?.getTime() === invalidatedAt.getTime() + ) + expect(invalidationUpdate).toEqual({ + toolCount: 0, + lastToolsRefresh: invalidatedAt, + }) + expect(cacheStore.has(serverKey)).toBe(false) + + releaseDiscoveryStatus?.([]) + await expect(discovery).rejects.toThrow('OAuth authorization required') + } finally { + vi.useRealTimers() + } + }) + + it('keeps a newer successful cache entry when an older failure finishes later', async () => { + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let rejectOlder: ((error: Error) => void) | undefined + const olderList = new Promise((_resolve, reject) => { + rejectOlder = reject + }) + mockListTools.mockReturnValueOnce(olderList).mockResolvedValueOnce([tool('new-tool', 'mcp-a')]) + + let releaseOlderFailureCache: (() => void) | undefined + const olderFailureCacheGate = new Promise((resolve) => { + releaseOlderFailureCache = resolve + }) + const defaultApply = mockCacheAdapter.applyMutationIfCurrent.getMockImplementation() + mockCacheAdapter.applyMutationIfCurrent.mockImplementationOnce( + async ( + scopeKey: string, + mutationId: number, + setEntry: { key: string; tools: unknown[]; ttlMs: number } | null, + deleteKeys: string[] + ) => { + await olderFailureCacheGate + return defaultApply?.(scopeKey, mutationId, setEntry, deleteKeys) ?? false + } + ) + + const older = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, false) + await vi.waitFor(() => expect(mockListTools).toHaveBeenCalledTimes(1)) + rejectOlder?.(new Error('Older request failed')) + + const newer = mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID, true) + await expect(newer).resolves.toEqual([tool('new-tool', 'mcp-a')]) + + releaseOlderFailureCache?.() + await expect(older).rejects.toThrow('Older request failed') + + expect(cacheStore.get(`workspace:${WORKSPACE_ID}:server:mcp-a`)?.tools).toEqual([ + tool('new-tool', 'mcp-a'), + ]) + expect(cacheStore.has(`workspace:${WORKSPACE_ID}:server:mcp-a:failure`)).toBe(false) + }) + it('retries a transient tools/list timeout and succeeds on the second attempt', async () => { mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) mockListTools @@ -484,6 +1514,59 @@ describe('McpService.discoverTools per-server caching', () => { expect(mockListTools).toHaveBeenCalledTimes(2) }) + it('reacquires mutation ownership after cache invalidation during retry backoff', async () => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + try { + const serverKey = `workspace:${WORKSPACE_ID}:server:mcp-a` + mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) + + let rejectFirstAttempt: ((error: Error) => void) | undefined + let markFirstAttemptStarted: (() => void) | undefined + const firstAttemptStarted = new Promise((resolve) => { + markFirstAttemptStarted = resolve + }) + mockListTools + .mockImplementationOnce( + () => + new Promise((_resolve, reject) => { + rejectFirstAttempt = reject + markFirstAttemptStarted?.() + }) + ) + .mockResolvedValueOnce([tool('retry-winner', 'mcp-a')]) + + const discovery = mcpService.discoverServerToolsWithMetadata( + USER_ID, + 'mcp-a', + WORKSPACE_ID, + true + ) + await firstAttemptStarted + + rejectFirstAttempt?.(new Error('Request timed out')) + await vi.advanceTimersByTimeAsync(0) + expect(mockListTools).toHaveBeenCalledTimes(1) + + await mcpService.clearCache(WORKSPACE_ID) + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(2) + + await vi.runAllTimersAsync() + await expect(discovery).resolves.toEqual({ + tools: [tool('retry-winner', 'mcp-a')], + state: 'published', + publicationOrder: expect.any(Date), + }) + + expect(mockCacheAdapter.beginMutation).toHaveBeenCalledTimes(3) + expect(cacheStore.get(serverKey)?.tools).toEqual([tool('retry-winner', 'mcp-a')]) + expect(mockUpdateSet).toHaveBeenCalledWith( + expect.objectContaining({ connectionStatus: 'connected', toolCount: 1 }) + ) + } finally { + vi.useRealTimers() + } + }) + it('persists and negative-caches per-server UnauthorizedError for headers auth', async () => { const reflectedCredential = 'Bearer static-secret-for-server-discovery' mockGetWorkspaceServersRows.mockResolvedValue([ @@ -507,7 +1590,9 @@ describe('McpService.discoverTools per-server caching', () => { }) ) expect(JSON.stringify(mockUpdateSet.mock.calls)).not.toContain(reflectedCredential) - expect(JSON.stringify(mockCacheAdapter.set.mock.calls)).not.toContain(reflectedCredential) + expect(JSON.stringify(mockCacheAdapter.applyMutationIfCurrent.mock.calls)).not.toContain( + reflectedCredential + ) expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) mockListTools.mockClear() @@ -531,10 +1616,11 @@ describe('McpService.discoverTools per-server caching', () => { lastError: null, }) ) - expect(mockCacheAdapter.set).not.toHaveBeenCalledWith( - `workspace:${WORKSPACE_ID}:server:mcp-a:failure`, - [], - expect.any(Number) + expect(mockCacheAdapter.applyMutationIfCurrent).not.toHaveBeenCalledWith( + expect.anything(), + expect.anything(), + expect.objectContaining({ key: `workspace:${WORKSPACE_ID}:server:mcp-a:failure` }), + expect.anything() ) mockResolveEnvVars.mockClear() @@ -564,6 +1650,45 @@ describe('McpService.discoverTools per-server caching', () => { ) }) + it('recomputes a failure count after a concurrent status update wins the CAS', async () => { + const beforeSuccess = dbRow('mcp-a', 'A', { + statusConfig: { consecutiveFailures: 2, lastSuccessfulDiscovery: null }, + }) + const afterSuccess = dbRow('mcp-a', 'A', { + statusConfig: { + consecutiveFailures: 0, + lastSuccessfulDiscovery: '2030-02-01T00:00:00.000Z', + }, + }) + mockGetWorkspaceServersRows + .mockResolvedValueOnce([beforeSuccess]) + .mockResolvedValueOnce([beforeSuccess]) + .mockResolvedValueOnce([afterSuccess]) + mockUpdateReturning.mockResolvedValueOnce([]).mockResolvedValueOnce([{ id: 'mcp-a' }]) + mockListTools.mockRejectedValueOnce(new Error('Connection refused')) + + await expect(mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID)).rejects.toThrow( + 'Connection refused' + ) + + const failureUpdates = mockUpdateSet.mock.calls + .map(([update]) => update) + .filter((update) => update.lastError === 'Connection failed') + expect(failureUpdates).toEqual([ + expect.objectContaining({ + connectionStatus: 'error', + statusConfig: { consecutiveFailures: 3, lastSuccessfulDiscovery: null }, + }), + expect.objectContaining({ + connectionStatus: 'disconnected', + statusConfig: { + consecutiveFailures: 1, + lastSuccessfulDiscovery: '2030-02-01T00:00:00.000Z', + }, + }), + ]) + }) + it('persists OAuth-required discovery as disconnected without a failure error', async () => { mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) mockListTools.mockRejectedValueOnce(new McpOauthAuthorizationRequiredError('mcp-a', 'A')) @@ -583,12 +1708,13 @@ describe('McpService.discoverTools per-server caching', () => { it('does not negative-cache a failure older than a successful discovery', async () => { mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) mockListTools.mockRejectedValueOnce(new Error('Older request failed')) - mockUpdateReturning.mockResolvedValueOnce([]) + mockUpdateReturning.mockResolvedValue([]) await expect(mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID)).rejects.toThrow( 'Older request failed' ) + mockUpdateReturning.mockResolvedValue([{ id: 'server-1' }]) mockListTools.mockResolvedValueOnce([tool('a1', 'mcp-a')]) const tools = await mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID) diff --git a/apps/sim/lib/mcp/service.ts b/apps/sim/lib/mcp/service.ts index 960007bf868..f23d8c66d51 100644 --- a/apps/sim/lib/mcp/service.ts +++ b/apps/sim/lib/mcp/service.ts @@ -1,3 +1,4 @@ +import { createHash } from 'node:crypto' import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js' import { StreamableHTTPError } from '@modelcontextprotocol/sdk/client/streamableHttp.js' import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js' @@ -7,7 +8,7 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { sleep } from '@sim/utils/helpers' import { backoffWithJitter } from '@sim/utils/retry' -import { and, eq, isNull, lte, or } from 'drizzle-orm' +import { and, eq, gte, isNull, lt, lte, or } from 'drizzle-orm' import { isTest } from '@/lib/core/config/env-flags' import { generateRequestId } from '@/lib/core/utils/request' import { McpClient } from '@/lib/mcp/client' @@ -17,6 +18,7 @@ import { validateMcpDomain, validateMcpServerSsrf, } from '@/lib/mcp/domain-check' +import { getMcpSafeErrorDiagnostics } from '@/lib/mcp/error-diagnostics' import { getOrCreateOauthRow, loadPreregisteredClient, @@ -27,6 +29,7 @@ import { resolveMcpConfigEnvVars } from '@/lib/mcp/resolve-config' import { createMcpCacheAdapter, getMcpCacheType, + type McpCacheMutationSet, type McpCacheStorageAdapter, } from '@/lib/mcp/storage' import { @@ -52,7 +55,38 @@ function failureCacheKey(workspaceId: string, serverId: string): string { return `workspace:${workspaceId}:server:${serverId}:failure` } +export function getTimestampMillisecondBounds(timestamp: string): { + startInclusive: Date + endExclusive: Date +} { + const startInclusive = new Date(timestamp) + return { + startInclusive, + endExclusive: new Date(startInclusive.getTime() + 1), + } +} + const FAILURE_CACHE_SENTINEL: McpTool[] = [] +const CACHE_MUTATION_BEGIN_ATTEMPTS = 2 +const CACHE_INVALIDATION_ATTEMPTS = 3 +const STATUS_UPDATE_CAS_ATTEMPTS = 3 +const STATUS_METADATA_RACE_RETRY_ATTEMPTS = 3 + +type CacheMutationResult = 'applied' | 'superseded' | 'unavailable' +type StatusPublicationResult = 'published' | Exclude + +export type McpServerDiscoveryState = + | 'cached' + | 'published' + | 'winner-cache' + | 'superseded' + | 'unavailable' + +export interface McpServerDiscoveryResult { + tools: McpTool[] + state: McpServerDiscoveryState + publicationOrder?: Date +} type DiscoveryOutcome = | { kind: 'cached'; tools: McpTool[] } @@ -60,17 +94,76 @@ type DiscoveryOutcome = kind: 'fetched' tools: McpTool[] resolvedConfig: McpServerConfig + sourceConfig: McpServerConfig resolvedIP: string | null + mutation: CacheMutation | null } - | { kind: 'oauth-pending' } + | { kind: 'oauth-pending'; config: McpServerConfig; mutation: CacheMutation | null } | { kind: 'unhealthy' } - // originalError preserves the type so markServerUnhealthy's instanceof - // exemption survives the getErrorMessage call. - | { kind: 'error'; message: string; originalError: unknown } + // originalError preserves the type so the OAuth exemption survives the + // getErrorMessage call. + | { + kind: 'error' + message: string + originalError: unknown + config: McpServerConfig + mutation: CacheMutation | null + } + +interface CacheMutation { + scopeKey: string + id: number +} + +interface DiscoveryRevisionInput { + transport: string | null + url: string | null + authType: string | null + oauthClientId: string | null + oauthClientSecret: string | null + headers: unknown + timeout: number | null + retries: number | null + enabled: boolean +} + +function getDiscoveryRevision(config: DiscoveryRevisionInput): string { + const headers = + config.headers && typeof config.headers === 'object' && !Array.isArray(config.headers) + ? Object.entries(config.headers as Record).sort(([left], [right]) => + left.localeCompare(right) + ) + : [] + return createHash('sha256') + .update( + JSON.stringify({ + transport: config.transport, + url: config.url, + authType: config.authType, + oauthClientId: config.oauthClientId, + oauthClientSecret: config.oauthClientSecret, + headers, + timeout: config.timeout, + retries: config.retries, + enabled: config.enabled, + }) + ) + .digest('hex') +} type ServerStatusUpdate = - | { outcome: 'connected'; toolCount: number } - | { outcome: 'failed'; error: string; discoveryStartedAt?: Date } + | { + outcome: 'connected' + toolCount: number + configUpdatedAt: string + publicationOrder: Date + } + | { + outcome: 'failed' + error: string + configUpdatedAt: string + publicationOrder: Date + } function isOauthAuthorizationError(error: unknown, authType: McpServerConfig['authType']): boolean { return ( @@ -79,18 +172,31 @@ function isOauthAuthorizationError(error: unknown, authType: McpServerConfig['au ) } -function getDiscoveryFailureMessage( - error: unknown, - authType: McpServerConfig['authType'], - fallback: string -): string { +function getDiscoveryFailureMessage(error: unknown, authType: McpServerConfig['authType']): string { if (authType !== 'oauth' && error instanceof UnauthorizedError) { return 'Authentication failed' } if (isTimeoutError(error)) { return 'The MCP server took too long to respond and timed out' } - return getErrorMessage(error, fallback) + if (error instanceof StreamableHTTPError) { + if (error.code === 401 || error.code === 403) return 'Authentication failed' + if (error.code === 429) return 'The MCP server is rate limited. Try again shortly.' + if (typeof error.code === 'number' && error.code >= 500) { + return 'The MCP server is temporarily unavailable' + } + } + const message = getErrorMessage(error, '').toLowerCase() + if ( + message.includes('econnrefused') || + message.includes('econnreset') || + message.includes('socket hang up') || + message.includes('fetch failed') || + message.includes('network') + ) { + return 'Unable to reach the MCP server' + } + return 'Connection failed' } function isTimeoutError(error: unknown): boolean { @@ -131,7 +237,7 @@ class McpService { private readonly cacheTimeout = MCP_CONSTANTS.CACHE_TIMEOUT private unsubscribeConnectionManager?: () => void // Keyed on (workspaceId, serverId, userId) — OAuth-scoped tokens vary per user. - private inflightServerDiscovery = new Map>() + private inflightServerDiscovery = new Map>() constructor() { this.cacheAdapter = createMcpCacheAdapter() @@ -139,19 +245,11 @@ class McpService { if (mcpConnectionManager) { this.unsubscribeConnectionManager = mcpConnectionManager.subscribe((event) => { - this.cacheAdapter - .delete(serverCacheKey(event.workspaceId, event.serverId)) - .catch((err) => - logger.warn(`Failed to invalidate cache for ${event.serverName} on listChanged:`, err) - ) - this.cacheAdapter - .delete(failureCacheKey(event.workspaceId, event.serverId)) - .catch((err) => - logger.warn( - `Failed to invalidate failure cache for ${event.serverName} on listChanged:`, - err - ) - ) + this.invalidateServerCache(event.workspaceId, event.serverId).catch((error) => { + logger.warn(`Failed to invalidate cache for ${event.serverName} on listChanged`, { + error: getMcpSafeErrorDiagnostics(error), + }) + }) }) } } @@ -214,6 +312,7 @@ class McpService { enabled: server.enabled, createdAt: server.createdAt.toISOString(), updatedAt: server.updatedAt.toISOString(), + discoveryRevision: getDiscoveryRevision(server), } } @@ -244,6 +343,7 @@ class McpService { enabled: server.enabled, createdAt: server.createdAt.toISOString(), updatedAt: server.updatedAt.toISOString(), + discoveryRevision: getDiscoveryRevision(server), })) .filter((config) => isMcpDomainAllowed(config.url)) } @@ -375,134 +475,157 @@ class McpService { ): Promise { try { const now = new Date() + const configUpdatedAt = getTimestampMillisecondBounds(update.configUpdatedAt) + const publicationConditions = and( + eq(mcpServers.id, serverId), + eq(mcpServers.workspaceId, workspaceId), + isNull(mcpServers.deletedAt), + gte(mcpServers.updatedAt, configUpdatedAt.startInclusive), + lt(mcpServers.updatedAt, configUpdatedAt.endExclusive), + or( + isNull(mcpServers.lastToolsRefresh), + lte(mcpServers.lastToolsRefresh, update.publicationOrder) + ) + ) if (update.outcome === 'connected') { - await db + const updatedServers = await db .update(mcpServers) .set({ connectionStatus: 'connected', lastConnected: now, lastError: null, toolCount: update.toolCount, - lastToolsRefresh: now, + // The cache mutation id is a monotonic millisecond timestamp. Use + // that same token here so cache and database publication cannot + // choose different winners during begin-mutation retries. + lastToolsRefresh: update.publicationOrder, statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: now.toISOString(), }, - updatedAt: now, }) - .where(eq(mcpServers.id, serverId)) - return true + .where(publicationConditions) + .returning({ id: mcpServers.id }) + return updatedServers.length > 0 } - const [currentServer] = await db - .select({ statusConfig: mcpServers.statusConfig }) - .from(mcpServers) - .where( - and( - eq(mcpServers.id, serverId), - eq(mcpServers.workspaceId, workspaceId), - isNull(mcpServers.deletedAt) + for (let attempt = 0; attempt < STATUS_UPDATE_CAS_ATTEMPTS; attempt++) { + const [currentServer] = await db + .select({ statusConfig: mcpServers.statusConfig }) + .from(mcpServers) + .where( + and( + eq(mcpServers.id, serverId), + eq(mcpServers.workspaceId, workspaceId), + isNull(mcpServers.deletedAt) + ) ) - ) - .limit(1) - - const storedConfig = currentServer?.statusConfig as Partial | null - const currentConfig: McpServerStatusConfig = { - consecutiveFailures: - typeof storedConfig?.consecutiveFailures === 'number' - ? storedConfig.consecutiveFailures - : 0, - lastSuccessfulDiscovery: storedConfig?.lastSuccessfulDiscovery ?? null, - } + .limit(1) + + if (!currentServer) return false + const storedConfig = currentServer.statusConfig as Partial | null + const currentConfig: McpServerStatusConfig = { + consecutiveFailures: + typeof storedConfig?.consecutiveFailures === 'number' + ? storedConfig.consecutiveFailures + : 0, + lastSuccessfulDiscovery: storedConfig?.lastSuccessfulDiscovery ?? null, + } + const newFailures = currentConfig.consecutiveFailures + 1 + const isErrorState = newFailures >= MCP_CONSTANTS.MAX_CONSECUTIVE_FAILURES + const statusConfigMatches = currentServer.statusConfig + ? eq(mcpServers.statusConfig, currentServer.statusConfig) + : isNull(mcpServers.statusConfig) - const newFailures = currentConfig.consecutiveFailures + 1 - const isErrorState = newFailures >= MCP_CONSTANTS.MAX_CONSECUTIVE_FAILURES + const updatedServers = await db + .update(mcpServers) + .set({ + connectionStatus: isErrorState ? 'error' : 'disconnected', + lastError: update.error || 'Unknown error', + toolCount: 0, + lastToolsRefresh: update.publicationOrder, + statusConfig: { + consecutiveFailures: newFailures, + lastSuccessfulDiscovery: currentConfig.lastSuccessfulDiscovery, + }, + }) + .where(and(publicationConditions, statusConfigMatches)) + .returning({ id: mcpServers.id }) - const updatedServers = await db - .update(mcpServers) - .set({ - connectionStatus: isErrorState ? 'error' : 'disconnected', - lastError: update.error || 'Unknown error', - statusConfig: { - consecutiveFailures: newFailures, - lastSuccessfulDiscovery: currentConfig.lastSuccessfulDiscovery, - }, - updatedAt: now, - }) - .where( - and( - eq(mcpServers.id, serverId), - eq(mcpServers.workspaceId, workspaceId), - isNull(mcpServers.deletedAt), - update.discoveryStartedAt - ? or( - isNull(mcpServers.lastConnected), - lte(mcpServers.lastConnected, update.discoveryStartedAt) - ) - : undefined + if (updatedServers.length === 0) continue + if (isErrorState) { + logger.warn( + `Server ${serverId} marked as error after ${newFailures} consecutive failures` ) - ) - .returning({ id: mcpServers.id }) - - if (isErrorState && updatedServers.length > 0) { - logger.warn(`Server ${serverId} marked as error after ${newFailures} consecutive failures`) + } + return true } - return updatedServers.length > 0 + return false } catch (err) { logger.error(`Failed to update server status for ${serverId}:`, err) return false } } - /** - * Negative-cache a discovery failure. OAuth-required errors are exempt so - * reconnects retry immediately. - */ - private async markServerUnhealthy( + private async applyServerCacheMutation( workspaceId: string, serverId: string, - error: unknown, - authType: McpServerConfig['authType'] - ): Promise { - if (isOauthAuthorizationError(error, authType)) { - return + mutation: CacheMutation | null, + setEntry: McpCacheMutationSet | null, + deleteKeys: string[] + ): Promise { + if (!mutation) { + // An unordered fallback cannot safely publish discovery state. An older + // ordered mutation could overwrite it and then lose the database CAS, + // while an unguarded delete could erase a newer publisher's result. + // Fail closed so cache and database publication remain one ordered unit. + return 'unavailable' } try { - await this.cacheAdapter.set( - failureCacheKey(workspaceId, serverId), - FAILURE_CACHE_SENTINEL, - MCP_CLIENT_CONSTANTS.FAILURE_CACHE_TTL_MS + const applied = await this.cacheAdapter.applyMutationIfCurrent( + mutation.scopeKey, + mutation.id, + setEntry, + deleteKeys ) - } catch (err) { - logger.warn(`Failed to write failure cache for server ${serverId}:`, err) + return applied ? 'applied' : 'superseded' + } catch (error) { + logger.warn(`Failed to atomically update cache for server ${serverId}`, { + workspaceId, + error: getMcpSafeErrorDiagnostics(error), + }) + return 'unavailable' } } private async markServerOauthPending( serverId: string, workspaceId: string, - discoveryStartedAt?: Date + configUpdatedAt: string, + publicationOrder: Date ): Promise { try { + const configUpdatedAtBounds = getTimestampMillisecondBounds(configUpdatedAt) const updatedServers = await db .update(mcpServers) .set({ connectionStatus: 'disconnected', lastError: null, - updatedAt: new Date(), + toolCount: 0, + lastToolsRefresh: publicationOrder, }) .where( and( eq(mcpServers.id, serverId), eq(mcpServers.workspaceId, workspaceId), isNull(mcpServers.deletedAt), - discoveryStartedAt - ? or( - isNull(mcpServers.lastConnected), - lte(mcpServers.lastConnected, discoveryStartedAt) - ) - : undefined + gte(mcpServers.updatedAt, configUpdatedAtBounds.startInclusive), + lt(mcpServers.updatedAt, configUpdatedAtBounds.endExclusive), + or( + isNull(mcpServers.lastToolsRefresh), + lte(mcpServers.lastToolsRefresh, publicationOrder) + ) ) ) .returning({ id: mcpServers.id }) @@ -513,6 +636,45 @@ class McpService { } } + /** + * Publish an invalidation's cache-order token without changing connection + * health. A list_changed notification comes from a live connection, while + * configuration lifecycle code persists its own intended connection state. + */ + private async markServerCacheInvalidated( + serverId: string, + workspaceId: string, + publicationOrder: Date + ): Promise { + try { + const updatedServers = await db + .update(mcpServers) + .set({ + toolCount: 0, + lastToolsRefresh: publicationOrder, + }) + .where( + and( + eq(mcpServers.id, serverId), + eq(mcpServers.workspaceId, workspaceId), + isNull(mcpServers.deletedAt), + or( + isNull(mcpServers.lastToolsRefresh), + lte(mcpServers.lastToolsRefresh, publicationOrder) + ) + ) + ) + .returning({ id: mcpServers.id }) + return updatedServers.length > 0 + } catch (error) { + logger.warn(`Failed to publish cache invalidation for server ${serverId}`, { + workspaceId, + error: getMcpSafeErrorDiagnostics(error), + }) + return false + } + } + private async isServerUnhealthy(workspaceId: string, serverId: string): Promise { try { const entry = await this.cacheAdapter.get(failureCacheKey(workspaceId, serverId)) @@ -522,12 +684,249 @@ class McpService { } } - private async clearServerFailure(workspaceId: string, serverId: string): Promise { + private async beginServerCacheMutation( + workspaceId: string, + serverId: string + ): Promise { + const scopeKey = serverCacheKey(workspaceId, serverId) + let lastError: unknown + for (let attempt = 0; attempt < CACHE_MUTATION_BEGIN_ATTEMPTS; attempt++) { + try { + return { scopeKey, id: await this.cacheAdapter.beginMutation(scopeKey) } + } catch (error) { + lastError = error + } + } + logger.warn(`Failed to order cache mutation for server ${serverId}`, { + attempts: CACHE_MUTATION_BEGIN_ATTEMPTS, + error: getMcpSafeErrorDiagnostics(lastError), + }) + return null + } + + private async invalidateServerCache(workspaceId: string, serverId: string): Promise { + const keys = [serverCacheKey(workspaceId, serverId), failureCacheKey(workspaceId, serverId)] + for (let attempt = 0; attempt < CACHE_INVALIDATION_ATTEMPTS; attempt++) { + // Reacquire a fresh ownership token after every unavailable transition. + // This orders the retry after work that completed during the unknown + // attempt, while work starting later can still supersede the new token. + const mutation = await this.beginServerCacheMutation(workspaceId, serverId) + if (!mutation) continue + + const cacheApplied = await this.applyServerCacheMutation( + workspaceId, + serverId, + mutation, + null, + keys + ) + if (cacheApplied === 'superseded') return + if (cacheApplied === 'unavailable') continue + + await this.markServerCacheInvalidated(serverId, workspaceId, new Date(mutation.id)) + return + } + + // Without an ownership-checked transition, deleting either key can erase + // a newer process's winner. Leave cache and database state unchanged. + logger.warn(`Cache invalidation unavailable for server ${serverId}`, { + workspaceId, + attempts: CACHE_INVALIDATION_ATTEMPTS, + }) + } + + private async getCurrentCachedTools( + workspaceId: string, + serverId: string + ): Promise { try { - await this.cacheAdapter.delete(failureCacheKey(workspaceId, serverId)) - } catch (err) { - logger.warn(`Failed to clear failure cache for server ${serverId}:`, err) + return (await this.cacheAdapter.get(serverCacheKey(workspaceId, serverId)))?.tools ?? null + } catch (error) { + logger.warn(`Failed to read current cache winner for server ${serverId}`, { + workspaceId, + error: getMcpSafeErrorDiagnostics(error), + }) + return null + } + } + + private async publishSuccessfulDiscovery( + workspaceId: string, + config: McpServerConfig, + mutation: CacheMutation | null, + tools: McpTool[] + ): Promise { + const cacheApplied = await this.applyServerCacheMutation( + workspaceId, + config.id, + mutation, + { + key: serverCacheKey(workspaceId, config.id), + tools, + ttlMs: this.cacheTimeout, + }, + [failureCacheKey(workspaceId, config.id)] + ) + if (cacheApplied !== 'applied') return cacheApplied + if (!mutation) return 'unavailable' + + const statusApplied = await this.updateServerStatus(config.id, workspaceId, { + outcome: 'connected', + toolCount: tools.length, + configUpdatedAt: config.updatedAt!, + publicationOrder: new Date(mutation.id), + }) + if (statusApplied) return 'published' + + const retryResult = await this.retryStatusPublicationAfterMetadataRace( + workspaceId, + config, + mutation, + (configUpdatedAt) => + this.updateServerStatus(config.id, workspaceId, { + outcome: 'connected', + toolCount: tools.length, + configUpdatedAt, + publicationOrder: new Date(mutation.id), + }) + ) + if (retryResult === 'superseded') { + await this.applyServerCacheMutation(workspaceId, config.id, mutation, null, [ + serverCacheKey(workspaceId, config.id), + failureCacheKey(workspaceId, config.id), + ]) } + return retryResult + } + + private async retryStatusPublicationAfterMetadataRace( + workspaceId: string, + config: McpServerConfig, + mutation: CacheMutation, + publishStatus: (configUpdatedAt: string) => Promise + ): Promise { + for (let attempt = 0; attempt < STATUS_METADATA_RACE_RETRY_ATTEMPTS; attempt++) { + const ownership = await this.applyServerCacheMutation( + workspaceId, + config.id, + mutation, + null, + [] + ) + if (ownership === 'superseded') return 'superseded' + if (ownership === 'unavailable') return 'unavailable' + + let currentConfig: McpServerConfig | null + try { + currentConfig = await this.getServerConfig(config.id, workspaceId) + } catch (error) { + logger.warn(`Failed to reread server ${config.id} for status publication`, { + workspaceId, + error: getMcpSafeErrorDiagnostics(error), + }) + return 'unavailable' + } + if ( + !currentConfig || + !config.discoveryRevision || + currentConfig.discoveryRevision !== config.discoveryRevision + ) { + return 'superseded' + } + if (await publishStatus(currentConfig.updatedAt!)) return 'published' + } + return 'unavailable' + } + + private async publishFailedDiscovery( + workspaceId: string, + config: McpServerConfig, + mutation: CacheMutation | null, + error: unknown, + message: string + ): Promise { + const cacheApplied = await this.applyServerCacheMutation( + workspaceId, + config.id, + mutation, + isOauthAuthorizationError(error, config.authType) + ? null + : { + key: failureCacheKey(workspaceId, config.id), + tools: FAILURE_CACHE_SENTINEL, + ttlMs: MCP_CLIENT_CONSTANTS.FAILURE_CACHE_TTL_MS, + }, + [serverCacheKey(workspaceId, config.id)] + ) + if (cacheApplied !== 'applied' || !mutation) return false + + const statusApplied = await this.updateServerStatus(config.id, workspaceId, { + outcome: 'failed', + error: message, + configUpdatedAt: config.updatedAt!, + publicationOrder: new Date(mutation.id), + }) + if (statusApplied) return true + + // A metadata-only edit can advance updatedAt without changing anything + // that affects discovery. Keep the failed publication only while this + // mutation still owns the cache and the connection configuration is + // unchanged, then retry the status CAS against the row's current token. + const retryResult = await this.retryStatusPublicationAfterMetadataRace( + workspaceId, + config, + mutation, + (configUpdatedAt) => + this.updateServerStatus(config.id, workspaceId, { + outcome: 'failed', + error: message, + configUpdatedAt, + publicationOrder: new Date(mutation.id), + }) + ) + if (retryResult === 'published') return true + + // Do not leave a negative-cache entry for a failure that lost the + // database publication CAS. + await this.applyServerCacheMutation(workspaceId, config.id, mutation, null, [ + failureCacheKey(workspaceId, config.id), + ]) + return false + } + + private async publishOauthPending( + workspaceId: string, + config: McpServerConfig, + mutation: CacheMutation | null + ): Promise { + const cacheApplied = await this.applyServerCacheMutation( + workspaceId, + config.id, + mutation, + null, + [serverCacheKey(workspaceId, config.id), failureCacheKey(workspaceId, config.id)] + ) + if (cacheApplied !== 'applied' || !mutation) return false + + const statusApplied = await this.markServerOauthPending( + config.id, + workspaceId, + config.updatedAt!, + new Date(mutation.id) + ) + if (statusApplied) return true + + // Metadata-only edits share discovery state with the original row. Retry + // only while this mutation still owns the cache and the discovery-relevant + // configuration has not changed. + const retryResult = await this.retryStatusPublicationAfterMetadataRace( + workspaceId, + config, + mutation, + (configUpdatedAt) => + this.markServerOauthPending(config.id, workspaceId, configUpdatedAt, new Date(mutation.id)) + ) + return retryResult === 'published' } async discoverTools( @@ -536,8 +935,6 @@ class McpService { forceRefresh = false ): Promise { const requestId = generateRequestId() - const discoveryStartedAt = new Date() - try { logger.info(`[${requestId}] Discovering MCP tools for workspace ${workspaceId}`) @@ -570,6 +967,8 @@ class McpService { } } + const mutation = await this.beginServerCacheMutation(workspaceId, config.id) + try { const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars( config, @@ -582,112 +981,107 @@ class McpService { logger.debug( `[${requestId}] Discovered ${tools.length} tools from server ${config.name}` ) - return { kind: 'fetched', tools, resolvedConfig, resolvedIP } + return { + kind: 'fetched', + tools, + resolvedConfig, + sourceConfig: config, + resolvedIP, + mutation, + } } finally { await client.disconnect() } } catch (error) { if (isOauthAuthorizationError(error, config.authType)) { - return { kind: 'oauth-pending' } + return { kind: 'oauth-pending', config, mutation } } return { kind: 'error', - message: getDiscoveryFailureMessage(error, config.authType, 'Unknown error'), + message: getDiscoveryFailureMessage(error, config.authType), originalError: error, + config, + mutation, } } }) ) - const allTools: McpTool[] = [] - const cacheWrites: Promise[] = [] - const deferredSideEffects: Promise[] = [] - const liveConnections: Array<{ - resolvedConfig: McpServerConfig - resolvedIP: string | null - }> = [] - let cachedCount = 0 - let fetchedCount = 0 - let failedCount = 0 - - outcomes.forEach((outcome, index) => { - const server = servers[index] - if (outcome.kind === 'cached') { - cachedCount++ - allTools.push(...outcome.tools) - return - } - if (outcome.kind === 'fetched') { - fetchedCount++ - allTools.push(...outcome.tools) - deferredSideEffects.push( - this.updateServerStatus(server.id, workspaceId, { - outcome: 'connected', - toolCount: outcome.tools.length, - }) - ) - cacheWrites.push( - this.cacheAdapter - .set(serverCacheKey(workspaceId, server.id), outcome.tools, this.cacheTimeout) - .catch((err) => - logger.warn(`[${requestId}] Cache write failed for ${server.name}:`, err) + const publications = await Promise.all( + outcomes.map(async (outcome, index) => { + const server = servers[index] + if (outcome.kind === 'cached') { + return { + tools: outcome.tools, + cached: 1, + fetched: 0, + failed: 0, + liveConnection: null, + } + } + if (outcome.kind === 'fetched') { + const publication = await this.publishSuccessfulDiscovery( + workspaceId, + outcome.sourceConfig, + outcome.mutation, + outcome.tools + ) + if (publication !== 'published') { + logger.info( + `[${requestId}] Discovery state was not published for server ${server.id}`, + { reason: publication } ) - ) - deferredSideEffects.push(this.clearServerFailure(workspaceId, server.id)) - liveConnections.push({ - resolvedConfig: outcome.resolvedConfig, - resolvedIP: outcome.resolvedIP, - }) - return - } - if (outcome.kind === 'oauth-pending') { - // Mark disconnected so the UI surfaces the re-auth button. - logger.info(`[${requestId}] Skipping server ${server.name}: OAuth authorization pending`) - deferredSideEffects.push( - this.markServerOauthPending(server.id, workspaceId, discoveryStartedAt).then( - () => undefined + } + const responseTools = + publication === 'published' || publication === 'unavailable' + ? outcome.tools + : ((await this.getCurrentCachedTools(workspaceId, server.id)) ?? []) + return { + tools: responseTools, + cached: 0, + fetched: 1, + failed: publication === 'superseded' && responseTools.length === 0 ? 1 : 0, + liveConnection: + publication === 'published' + ? { + resolvedConfig: outcome.resolvedConfig, + resolvedIP: outcome.resolvedIP, + } + : null, + } + } + if (outcome.kind === 'oauth-pending') { + logger.info( + `[${requestId}] Skipping server ${server.name}: OAuth authorization pending` ) - ) - return - } - if (outcome.kind === 'unhealthy') { - // Status was persisted on the original failure; nothing to re-write. - failedCount++ - return - } - failedCount++ - logger.warn( - `[${requestId}] Failed to discover tools from server ${server.name}: ${outcome.message}` - ) - deferredSideEffects.push( - this.updateServerStatus(server.id, workspaceId, { - outcome: 'failed', + await this.publishOauthPending(workspaceId, outcome.config, outcome.mutation) + return { tools: [], cached: 0, fetched: 0, failed: 0, liveConnection: null } + } + if (outcome.kind === 'unhealthy') { + return { tools: [], cached: 0, fetched: 0, failed: 1, liveConnection: null } + } + + logger.warn(`[${requestId}] Failed to discover tools from server ${server.name}`, { error: outcome.message, - discoveryStartedAt, - }).then(async (statusApplied) => { - if (!statusApplied) return - await Promise.allSettled([ - this.markServerUnhealthy( - workspaceId, - server.id, - outcome.originalError, - server.authType - ), - this.cacheAdapter - .delete(serverCacheKey(workspaceId, server.id)) - .catch((err) => - logger.warn(`[${requestId}] Cache delete failed for ${server.name}:`, err) - ), - ]) }) - ) - }) + await this.publishFailedDiscovery( + workspaceId, + outcome.config, + outcome.mutation, + outcome.originalError, + outcome.message + ) + return { tools: [], cached: 0, fetched: 0, failed: 1, liveConnection: null } + }) + ) - // Await cache writes so a follow-up discoverTools sees consistent state. - await Promise.allSettled(cacheWrites) - // Each deferred side-effect self-logs failures, so we just mark the - // promises as handled to avoid unhandled-rejection warnings. - for (const p of deferredSideEffects) p.catch(() => {}) + const allTools = publications.flatMap((publication) => publication.tools) + const cachedCount = publications.reduce((sum, publication) => sum + publication.cached, 0) + const fetchedCount = publications.reduce((sum, publication) => sum + publication.fetched, 0) + const failedCount = publications.reduce((sum, publication) => sum + publication.failed, 0) + const liveConnections = publications.flatMap((publication) => + publication.liveConnection ? [publication.liveConnection] : [] + ) if (mcpConnectionManager) { for (const conn of liveConnections) { @@ -725,6 +1119,16 @@ class McpService { workspaceId: string, forceRefresh = false ): Promise { + return (await this.discoverServerToolsWithMetadata(userId, serverId, workspaceId, forceRefresh)) + .tools + } + + async discoverServerToolsWithMetadata( + userId: string, + serverId: string, + workspaceId: string, + forceRefresh = false + ): Promise { const inflightKey = `${workspaceId}:${serverId}:${userId}:${forceRefresh ? 'force' : 'cache'}` const existing = this.inflightServerDiscovery.get(inflightKey) if (existing) return existing @@ -746,9 +1150,8 @@ class McpService { serverId: string, workspaceId: string, forceRefresh: boolean - ): Promise { + ): Promise { const requestId = generateRequestId() - const discoveryStartedAt = new Date() const maxRetries = 2 if (!forceRefresh) { @@ -756,7 +1159,7 @@ class McpService { const cached = await this.cacheAdapter.get(serverCacheKey(workspaceId, serverId)) if (cached) { logger.debug(`[${requestId}] Cache hit for server ${serverId}`) - return cached.tools + return { tools: cached.tools, state: 'cached' } } } catch (error) { logger.warn(`[${requestId}] Cache read failed for server ${serverId}:`, error) @@ -771,17 +1174,20 @@ class McpService { } for (let attempt = 0; attempt < maxRetries; attempt++) { - let authType: McpServerConfig['authType'] + let config: McpServerConfig | null = null + // Begin a fresh mutation per attempt. A retry that succeeds after a + // concurrent clearCache must publish under a current ownership id — a + // stale pre-loop id would lose the CAS and drop otherwise-valid tools. + const mutation = await this.beginServerCacheMutation(workspaceId, serverId) try { logger.info( `[${requestId}] Discovering tools from server ${serverId} for user ${userId}${attempt > 0 ? ` (attempt ${attempt + 1})` : ''}` ) - const config = await this.getServerConfig(serverId, workspaceId) + config = await this.getServerConfig(serverId, workspaceId) if (!config) { throw new Error(`Server ${serverId} not found or not accessible`) } - authType = config.authType const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars( config, @@ -793,48 +1199,50 @@ class McpService { try { const tools = await client.listTools() logger.info(`[${requestId}] Discovered ${tools.length} tools from server ${config.name}`) - 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.clearServerFailure(workspaceId, serverId), - this.updateServerStatus(serverId, workspaceId, { - outcome: 'connected', - toolCount: tools.length, - }), - ]) - return tools + const publication = await this.publishSuccessfulDiscovery( + workspaceId, + config, + mutation, + tools + ) + if (publication !== 'published') { + logger.info(`[${requestId}] Discovery state was not published for server ${serverId}`, { + reason: publication, + }) + if (publication === 'superseded') { + const winner = await this.getCurrentCachedTools(workspaceId, serverId) + return winner + ? { tools: winner, state: 'winner-cache' } + : { tools: [], state: 'superseded' } + } + return { tools, state: 'unavailable' } + } + if (!mutation) return { tools, state: 'unavailable' } + return { tools, state: 'published', publicationOrder: new Date(mutation.id) } } finally { await client.disconnect() } } catch (error) { if (isRetryableDiscoveryError(error) && attempt < maxRetries - 1) { logger.warn( - `[${requestId}] Transient error discovering tools from server ${serverId}, retrying (attempt ${attempt + 1}):`, - error + `[${requestId}] Transient error discovering tools from server ${serverId}, retrying (attempt ${attempt + 1})`, + { error: getMcpSafeErrorDiagnostics(error) } ) await sleep(backoffWithJitter(attempt + 1, null, { baseMs: 250, maxMs: 2000 })) continue } - // Drop positive cache so a follow-up doesn't return stale tools. - const statusApplied = isOauthAuthorizationError(error, authType) - ? await this.markServerOauthPending(serverId, workspaceId, discoveryStartedAt) - : await this.updateServerStatus(serverId, workspaceId, { - outcome: 'failed', - error: getDiscoveryFailureMessage(error, authType, 'Connection failed'), - discoveryStartedAt, - }) - if (statusApplied) { - await Promise.allSettled([ - this.cacheAdapter - .delete(serverCacheKey(workspaceId, serverId)) - .catch((err) => - logger.warn(`[${requestId}] Cache delete failed for ${serverId}:`, err) - ), - this.markServerUnhealthy(workspaceId, serverId, error, authType), - ]) + if (config) { + if (isOauthAuthorizationError(error, config.authType)) { + await this.publishOauthPending(workspaceId, config, mutation) + } else { + await this.publishFailedDiscovery( + workspaceId, + config, + mutation, + error, + getDiscoveryFailureMessage(error, config.authType) + ) + } } throw error } @@ -895,7 +1303,7 @@ class McpService { status: 'error', toolCount: 0, lastSeen: undefined, - error: getDiscoveryFailureMessage(error, config.authType, 'Connection failed'), + error: getDiscoveryFailureMessage(error, config.authType), }) } } @@ -917,12 +1325,7 @@ class McpService { .select({ id: mcpServers.id }) .from(mcpServers) .where(eq(mcpServers.workspaceId, workspaceId)) - await Promise.allSettled( - rows.flatMap((r) => [ - this.cacheAdapter.delete(serverCacheKey(workspaceId, r.id)), - this.cacheAdapter.delete(failureCacheKey(workspaceId, r.id)), - ]) - ) + await Promise.allSettled(rows.map((row) => this.invalidateServerCache(workspaceId, row.id))) logger.debug(`Cleared MCP tool cache for workspace ${workspaceId} (${rows.length} servers)`) } else { await this.cacheAdapter.clear() diff --git a/apps/sim/lib/mcp/storage/adapter.ts b/apps/sim/lib/mcp/storage/adapter.ts index 9332bb371b5..87b3c2fe644 100644 --- a/apps/sim/lib/mcp/storage/adapter.ts +++ b/apps/sim/lib/mcp/storage/adapter.ts @@ -5,10 +5,30 @@ export interface McpCacheEntry { expiry: number // Unix timestamp ms } +export interface McpCacheMutationSet { + key: string + tools: McpTool[] + ttlMs: number +} + export interface McpCacheStorageAdapter { get(key: string): Promise set(key: string, tools: McpTool[], ttlMs: number): Promise delete(key: string): Promise + /** + * Starts an ordered mutation for one server and returns a monotonic Unix + * timestamp in milliseconds. Conditional writes using an older mutation id + * must be ignored so a slow discovery cannot overwrite a newer result. The + * same value orders database publication for an end-to-end consistent state. + */ + beginMutation(scopeKey: string): Promise + /** Atomically applies one server's complete cache state if this mutation still owns it. */ + applyMutationIfCurrent( + scopeKey: string, + mutationId: number, + setEntry: McpCacheMutationSet | null, + deleteKeys: string[] + ): Promise clear(): Promise dispose(): void } diff --git a/apps/sim/lib/mcp/storage/index.ts b/apps/sim/lib/mcp/storage/index.ts index 0990173c82d..f361d5d2548 100644 --- a/apps/sim/lib/mcp/storage/index.ts +++ b/apps/sim/lib/mcp/storage/index.ts @@ -1,2 +1,2 @@ -export type { McpCacheStorageAdapter } from './adapter' +export type { McpCacheMutationSet, McpCacheStorageAdapter } from './adapter' export { createMcpCacheAdapter, getMcpCacheType } from './factory' diff --git a/apps/sim/lib/mcp/storage/memory-cache.test.ts b/apps/sim/lib/mcp/storage/memory-cache.test.ts index c97944288fd..94e5c80e28d 100644 --- a/apps/sim/lib/mcp/storage/memory-cache.test.ts +++ b/apps/sim/lib/mcp/storage/memory-cache.test.ts @@ -165,6 +165,74 @@ describe('MemoryMcpCache', () => { }) }) + describe('ordered mutations', () => { + it('prevents an older discovery from overwriting a newer cache result', async () => { + const beforeMutation = Date.now() + const older = await cache.beginMutation('server-1') + const newer = await cache.beginMutation('server-1') + + expect(older).toBeGreaterThanOrEqual(beforeMutation) + expect(newer).toBeGreaterThan(older) + + expect( + await cache.applyMutationIfCurrent( + 'server-1', + newer, + { key: 'server-1:tools', tools: [createTool('new-tool')], ttlMs: 60000 }, + [] + ) + ).toBe(true) + expect( + await cache.applyMutationIfCurrent( + 'server-1', + older, + { key: 'server-1:failure', tools: [], ttlMs: 60000 }, + [] + ) + ).toBe(false) + + expect((await cache.get('server-1:tools'))?.tools).toEqual([createTool('new-tool')]) + expect(await cache.get('server-1:failure')).toBeNull() + }) + + it('atomically replaces tools and failure state for one mutation', async () => { + await cache.set('server-1:failure', [], 60000) + const mutation = await cache.beginMutation('server-1') + + expect( + await cache.applyMutationIfCurrent( + 'server-1', + mutation, + { + key: 'server-1:tools', + tools: [createTool('new-tool')], + ttlMs: 60000, + }, + ['server-1:failure'] + ) + ).toBe(true) + + expect((await cache.get('server-1:tools'))?.tools).toEqual([createTool('new-tool')]) + expect(await cache.get('server-1:failure')).toBeNull() + }) + + it('invalidates in-flight mutations when the cache is cleared', async () => { + const mutation = await cache.beginMutation('server-1') + + await cache.clear() + + expect( + await cache.applyMutationIfCurrent( + 'server-1', + mutation, + { key: 'server-1:tools', tools: [createTool('stale-tool')], ttlMs: 60000 }, + [] + ) + ).toBe(false) + expect(await cache.get('server-1:tools')).toBeNull() + }) + }) + describe('clear', () => { it('removes all entries from cache', async () => { const tools = [createTool('tool-1')] diff --git a/apps/sim/lib/mcp/storage/memory-cache.ts b/apps/sim/lib/mcp/storage/memory-cache.ts index b9d54194827..ac2c56176d9 100644 --- a/apps/sim/lib/mcp/storage/memory-cache.ts +++ b/apps/sim/lib/mcp/storage/memory-cache.ts @@ -1,12 +1,14 @@ import { createLogger } from '@sim/logger' import type { McpTool } from '@/lib/mcp/types' import { MCP_CONSTANTS } from '@/lib/mcp/utils' -import type { McpCacheEntry, McpCacheStorageAdapter } from './adapter' +import type { McpCacheEntry, McpCacheMutationSet, McpCacheStorageAdapter } from './adapter' const logger = createLogger('McpMemoryCache') export class MemoryMcpCache implements McpCacheStorageAdapter { private cache = new Map() + private mutationVersions = new Map() + private nextMutationId = 0 private readonly maxCacheSize = MCP_CONSTANTS.MAX_CACHE_SIZE private cleanupInterval: NodeJS.Timeout | null = null @@ -88,7 +90,38 @@ export class MemoryMcpCache implements McpCacheStorageAdapter { this.cache.delete(key) } + async beginMutation(scopeKey: string): Promise { + const mutationId = Math.max(this.nextMutationId + 1, Date.now()) + this.nextMutationId = mutationId + this.mutationVersions.set(scopeKey, mutationId) + return mutationId + } + + async applyMutationIfCurrent( + scopeKey: string, + mutationId: number, + setEntry: McpCacheMutationSet | null, + deleteKeys: string[] + ): Promise { + if (this.mutationVersions.get(scopeKey) !== mutationId) return false + + if (setEntry) { + this.cache.set(setEntry.key, { + tools: setEntry.tools, + expiry: Date.now() + setEntry.ttlMs, + }) + } + for (const key of deleteKeys) this.cache.delete(key) + this.evictIfNeeded() + return true + } + async clear(): Promise { + for (const scopeKey of this.mutationVersions.keys()) { + const mutationId = Math.max(this.nextMutationId + 1, Date.now()) + this.nextMutationId = mutationId + this.mutationVersions.set(scopeKey, mutationId) + } this.cache.clear() } @@ -98,6 +131,7 @@ export class MemoryMcpCache implements McpCacheStorageAdapter { this.cleanupInterval = null } this.cache.clear() + this.mutationVersions.clear() logger.info('Memory cache disposed') } } diff --git a/apps/sim/lib/mcp/storage/redis-cache.test.ts b/apps/sim/lib/mcp/storage/redis-cache.test.ts new file mode 100644 index 00000000000..af158a3a43c --- /dev/null +++ b/apps/sim/lib/mcp/storage/redis-cache.test.ts @@ -0,0 +1,107 @@ +/** + * @vitest-environment node + */ +import type Redis from 'ioredis' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { RedisMcpCache } from '@/lib/mcp/storage/redis-cache' +import type { McpTool } from '@/lib/mcp/types' + +const tool: McpTool = { + name: 'new-tool', + description: 'New tool', + inputSchema: { type: 'object' }, + serverId: 'server-1', + serverName: 'Server 1', +} + +describe('RedisMcpCache ordered mutations', () => { + const multi = { + incr: vi.fn(), + pexpire: vi.fn(), + exec: vi.fn(), + } + const redis = { + multi: vi.fn(() => multi), + eval: vi.fn(), + scan: vi.fn(), + del: vi.fn(), + } + const cache = new RedisMcpCache(redis as unknown as Redis) + + beforeEach(() => { + vi.clearAllMocks() + redis.eval.mockReset() + redis.scan.mockReset() + redis.del.mockReset() + multi.incr.mockReturnValue(multi) + multi.pexpire.mockReturnValue(multi) + multi.exec.mockResolvedValue([ + [null, 7], + [null, 1], + ]) + }) + + it('allocates a timestamp-based per-server mutation id with an expiry', async () => { + const mutationId = 1_900_000_000_000 + redis.eval.mockResolvedValueOnce(mutationId) + + await expect(cache.beginMutation('workspace:w:server:s')).resolves.toBe(mutationId) + + expect(redis.eval).toHaveBeenCalledWith( + expect.stringMatching(/redis\.call\('TIME'\).*math\.max/s), + 1, + 'mcp:tools-mutation:workspace:w:server:s', + String(24 * 60 * 60 * 1000) + ) + }) + + it('atomically replaces the complete cache state for the current mutation', async () => { + redis.eval.mockResolvedValueOnce(1) + + await expect( + cache.applyMutationIfCurrent( + 'workspace:w:server:s', + 7, + { + key: 'workspace:w:server:s', + tools: [tool], + ttlMs: 60_000, + }, + ['workspace:w:server:s:failure'] + ) + ).resolves.toBe(true) + + expect(redis.eval).toHaveBeenCalledWith( + expect.stringMatching(/redis\.call\('SET'.*redis\.call\('DEL'/s), + 3, + 'mcp:tools-mutation:workspace:w:server:s', + 'mcp:tools:workspace:w:server:s', + 'mcp:tools:workspace:w:server:s:failure', + '7', + '1', + expect.stringContaining('new-tool'), + '60000' + ) + }) + + it('invalidates mutation owners before deleting entries during a full clear', async () => { + redis.scan + .mockResolvedValueOnce(['0', ['mcp:tools-mutation:workspace:w:server:s']]) + .mockResolvedValueOnce([ + '0', + ['mcp:tools:workspace:w:server:s', 'mcp:tools:workspace:w:server:s:failure'], + ]) + redis.del.mockResolvedValueOnce(2) + + await cache.clear() + + expect(multi.incr).toHaveBeenCalledWith('mcp:tools-mutation:workspace:w:server:s') + expect(redis.del).toHaveBeenCalledWith( + 'mcp:tools:workspace:w:server:s', + 'mcp:tools:workspace:w:server:s:failure' + ) + expect(multi.exec.mock.invocationCallOrder[0]).toBeLessThan( + redis.del.mock.invocationCallOrder[0] + ) + }) +}) diff --git a/apps/sim/lib/mcp/storage/redis-cache.ts b/apps/sim/lib/mcp/storage/redis-cache.ts index f04b9fea119..04d5b95f493 100644 --- a/apps/sim/lib/mcp/storage/redis-cache.ts +++ b/apps/sim/lib/mcp/storage/redis-cache.ts @@ -1,11 +1,35 @@ import { createLogger } from '@sim/logger' import type Redis from 'ioredis' import type { McpTool } from '@/lib/mcp/types' -import type { McpCacheEntry, McpCacheStorageAdapter } from './adapter' +import type { McpCacheEntry, McpCacheMutationSet, McpCacheStorageAdapter } from './adapter' const logger = createLogger('McpRedisCache') const REDIS_KEY_PREFIX = 'mcp:tools:' +const MUTATION_KEY_PREFIX = 'mcp:tools-mutation:' +const MUTATION_TTL_MS = 24 * 60 * 60 * 1000 + +const BEGIN_MUTATION = ` +local current = tonumber(redis.call('GET', KEYS[1]) or '0') +local redisTime = redis.call('TIME') +local now = tonumber(redisTime[1]) * 1000 + math.floor(tonumber(redisTime[2]) / 1000) +local mutationId = math.max(current + 1, now) +redis.call('SET', KEYS[1], tostring(mutationId), 'PX', ARGV[1]) +return mutationId +` + +const APPLY_MUTATION_IF_CURRENT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return 0 +end +if ARGV[2] == '1' then + redis.call('SET', KEYS[2], ARGV[3], 'PX', ARGV[4]) +end +for index = 3, #KEYS do + redis.call('DEL', KEYS[index]) +end +return 1 +` export class RedisMcpCache implements McpCacheStorageAdapter { constructor(private redis: Redis) {} @@ -14,6 +38,10 @@ export class RedisMcpCache implements McpCacheStorageAdapter { return `${REDIS_KEY_PREFIX}${key}` } + private getMutationKey(scopeKey: string): string { + return `${MUTATION_KEY_PREFIX}${scopeKey}` + } + async get(key: string): Promise { try { const redisKey = this.getKey(key) @@ -61,11 +89,86 @@ export class RedisMcpCache implements McpCacheStorageAdapter { } } + async beginMutation(scopeKey: string): Promise { + try { + const mutationKey = this.getMutationKey(scopeKey) + const mutationId = await this.redis.eval( + BEGIN_MUTATION, + 1, + mutationKey, + String(MUTATION_TTL_MS) + ) + if (typeof mutationId !== 'number') { + throw new Error('Redis did not return an MCP cache mutation id') + } + return mutationId + } catch (error) { + logger.error('Redis cache mutation start error:', error) + throw error + } + } + + async applyMutationIfCurrent( + scopeKey: string, + mutationId: number, + setEntry: McpCacheMutationSet | null, + deleteKeys: string[] + ): Promise { + try { + const entry = setEntry + ? JSON.stringify({ + tools: setEntry.tools, + expiry: Date.now() + setEntry.ttlMs, + } satisfies McpCacheEntry) + : '' + const keys = [ + this.getMutationKey(scopeKey), + setEntry ? this.getKey(setEntry.key) : this.getMutationKey(scopeKey), + ...deleteKeys.map((key) => this.getKey(key)), + ] + const result = await this.redis.eval( + APPLY_MUTATION_IF_CURRENT, + keys.length, + ...keys, + String(mutationId), + setEntry ? '1' : '0', + entry, + String(setEntry?.ttlMs ?? 0) + ) + return result === 1 + } catch (error) { + logger.error('Redis atomic cache mutation error:', error) + throw error + } + } + async clear(): Promise { try { let cursor = '0' - let deletedCount = 0 + // Invalidate existing mutation owners before deleting their cache + // entries. An old writer either commits before this point and is then + // deleted, or observes the advanced id and cannot commit afterward. + do { + const [nextCursor, keys] = await this.redis.scan( + cursor, + 'MATCH', + `${MUTATION_KEY_PREFIX}*`, + 'COUNT', + 100 + ) + cursor = nextCursor + if (keys.length > 0) { + const transaction = this.redis.multi() + for (const key of keys) { + transaction.incr(key) + transaction.pexpire(key, MUTATION_TTL_MS) + } + await transaction.exec() + } + } while (cursor !== '0') + cursor = '0' + let deletedCount = 0 do { const [nextCursor, keys] = await this.redis.scan( cursor, diff --git a/apps/sim/lib/mcp/types.ts b/apps/sim/lib/mcp/types.ts index 575e86ea2f8..dd8d4dc51d6 100644 --- a/apps/sim/lib/mcp/types.ts +++ b/apps/sim/lib/mcp/types.ts @@ -30,6 +30,8 @@ export interface McpServerConfig { statusConfig?: McpServerStatusConfig createdAt?: string updatedAt?: string + /** Internal hash of fields that affect discovery; excludes display-only metadata. */ + discoveryRevision?: string } export interface McpVersionInfo {