diff --git a/apps/sim/app/api/mcp/oauth/callback/route.ts b/apps/sim/app/api/mcp/oauth/callback/route.ts index 7694f23b010..79b7058ddea 100644 --- a/apps/sim/app/api/mcp/oauth/callback/route.ts +++ b/apps/sim/app/api/mcp/oauth/callback/route.ts @@ -125,12 +125,19 @@ export const GET = withRouteHandler(async (request: NextRequest) => { serverId?: string ) => htmlClose(message, ok, reason, serverId, state) - const initialRow = state ? await loadOauthRowByState(state).catch(() => null) : null + const initialRow = state + ? await timedStep('loadOauthRowByState', 15_000, () => loadOauthRowByState(state)).catch( + () => null + ) + : null const stateRowServerId = initialRow?.mcpServerId if (errorParam) { logger.warn(`MCP OAuth callback received error: ${errorParam}`) - if (initialRow) await clearState(initialRow.id, 'callback:provider_error').catch(() => {}) + if (initialRow) + await timedStep('clearState(provider_error)', 10_000, () => + clearState(initialRow.id, 'callback:provider_error') + ).catch(() => {}) return respond(`Authorization failed: ${errorParam}`, false, 'provider_error', stateRowServerId) } if (!state || !code) { @@ -144,7 +151,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => { let serverId: string | undefined try { - const session = await getSession() + const session = await timedStep('getSession', 15_000, () => getSession()) if (!session?.user?.id) { return respond( 'You must be signed in to complete authorization.', @@ -169,11 +176,13 @@ export const GET = withRouteHandler(async (request: NextRequest) => { ) } - const [server] = await db - .select({ id: mcpServers.id, url: mcpServers.url, workspaceId: mcpServers.workspaceId }) - .from(mcpServers) - .where(and(eq(mcpServers.id, row.mcpServerId), isNull(mcpServers.deletedAt))) - .limit(1) + const [server] = await timedStep('loadServer', 15_000, () => + db + .select({ id: mcpServers.id, url: mcpServers.url, workspaceId: mcpServers.workspaceId }) + .from(mcpServers) + .where(and(eq(mcpServers.id, row.mcpServerId), isNull(mcpServers.deletedAt))) + .limit(1) + ) if (!server || !server.url) { return respond('Server no longer exists.', false, 'server_gone', serverId) } diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx index a4a937a5e27..3cfcb3fe3e8 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx @@ -79,7 +79,6 @@ import { useCreateMcpServer, useForceRefreshMcpTools, useMcpServers, - useMcpToolsEvents, useStoredMcpTools, } from '@/hooks/queries/mcp' import { useWorkflowState, useWorkflows } from '@/hooks/queries/workflows' @@ -570,7 +569,6 @@ export const ToolInput = memo(function ToolInput({ const { data: mcpServers = [], isLoading: mcpServersLoading } = useMcpServers(workspaceId) const { data: storedMcpTools = [] } = useStoredMcpTools(workspaceId) const forceRefreshMcpTools = useForceRefreshMcpTools().mutate - useMcpToolsEvents(workspaceId) const { navigateToSettings } = useSettingsNavigation() const createMcpServer = useCreateMcpServer() const { startOauthForServer } = useMcpOauthPopup({ workspaceId }) diff --git a/apps/sim/hooks/queries/mcp.test.tsx b/apps/sim/hooks/queries/mcp.test.tsx index ad6c120f516..18e3f2d5e34 100644 --- a/apps/sim/hooks/queries/mcp.test.tsx +++ b/apps/sim/hooks/queries/mcp.test.tsx @@ -103,13 +103,29 @@ function mockServers(servers: McpServer[]) { }) } +// jsdom has no EventSource; useMcpToolsQuery mounts the shared SSE subscription. +class FakeEventSource { + onopen: (() => void) | null = null + onerror: (() => void) | null = null + constructor(public url: string) {} + addEventListener(): void {} + close(): void {} +} + describe('useMcpToolsQuery', () => { beforeEach(() => { vi.clearAllMocks() + ;(globalThis as unknown as { EventSource: unknown }).EventSource = FakeEventSource }) afterEach(() => { vi.restoreAllMocks() + // mcp.ts captured these Map/Set instances in module consts at import, so reassigning the + // globalThis property wouldn't reset what the module uses — clear the shared instances. + ;( + globalThis as unknown as { __mcp_sse_connections?: Map } + ).__mcp_sse_connections?.clear() + ;(globalThis as unknown as { __mcp_sse_subscribed?: Set }).__mcp_sse_subscribed?.clear() }) it('does not auto-discover disconnected or errored OAuth servers', async () => { @@ -139,7 +155,11 @@ describe('useMcpToolsQuery', () => { const { unmount } = renderHookWithClient(() => useMcpToolsQuery(WORKSPACE_ID)) await flush() - expect(mockRequestJson).toHaveBeenCalledTimes(3) + // Both eligible servers get discovered. + const discoveryCalls = mockRequestJson.mock.calls.filter( + ([contract]) => contract === discoverMcpToolsContract + ) + expect(discoveryCalls).toHaveLength(2) expect(mockRequestJson).toHaveBeenCalledWith( discoverMcpToolsContract, expect.objectContaining({ diff --git a/apps/sim/hooks/queries/mcp.ts b/apps/sim/hooks/queries/mcp.ts index 8daab301235..f7fbffe9176 100644 --- a/apps/sim/hooks/queries/mcp.ts +++ b/apps/sim/hooks/queries/mcp.ts @@ -42,7 +42,13 @@ const logger = createLogger('McpQueries') export type { McpServerStatusConfig, McpTool, StoredMcpTool } export const MCP_SERVER_LIST_STALE_TIME = 60 * 1000 -export const MCP_SERVER_TOOLS_STALE_TIME = 30 * 1000 +/** + * Tool discovery is kept fresh by the `list_changed` → SSE push (see `useMcpToolsEvents`), + * so the query only needs a re-probe-on-visit fallback for servers without push. Matches the + * server-side cache TTL (`MCP_CONSTANTS.CACHE_TIMEOUT`) — no reference MCP client re-probes + * more often than its cache; real changes arrive via push regardless of this value. + */ +export const MCP_SERVER_TOOLS_STALE_TIME = 5 * 60 * 1000 export const MCP_STORED_TOOL_LIST_STALE_TIME = 60 * 1000 export const MCP_ALLOWED_DOMAINS_STALE_TIME = 5 * 60 * 1000 @@ -140,6 +146,11 @@ function isServerEligibleForDiscovery(server: McpServer, workspaceId: string): b export function useMcpToolsQuery(workspaceId: string) { const queryClient = useQueryClient() const { data: servers, isLoading: serversLoading } = useMcpServers(workspaceId) + // Push is intrinsic to consuming the tools query: every surface that reads tools (settings, + // tool picker, dynamic args, tool selector, canvas block) gets real-time `list_changed` + // refresh via the shared, reference-counted subscription — so the 5-min stale time is always + // push-backed and no consumer is left re-probing. + useMcpToolsEvents(workspaceId) /** * Skip disabled rows, rows retained from a previous workspace, and OAuth rows @@ -192,10 +203,10 @@ export function useMcpToolsQuery(workspaceId: string) { const serverId = serverIds[index] const status = serverId ? statusById.get(serverId) : undefined const persistentlyFailed = status === 'error' || status === 'disconnected' - // Keep last-known-good tools through a transient failure (React Query retains `data`, the - // stored status is still healthy) so a populated server doesn't blank — but drop them once - // the stored status crosses its failure threshold, so the workflow editor stops offering a - // dead server's stale tools. Matches how reference MCP clients treat transient vs. closed. + // Keep last-known-good tools while the stored status is still `connected` (React Query + // retains `data` across a failed refetch, so a populated server doesn't blank on a + // transient probe error) — but drop them once the stored status leaves `connected` + // (disconnected/error), so the workflow editor stops offering a dead server's stale tools. if (result.data && (!result.isError || !persistentlyFailed)) { tools.push(...result.data) hasData = true @@ -527,6 +538,12 @@ const sseConnections: Map = ((globalThis as Record)[SSE_KEY] as Map) ?? ((globalThis as Record)[SSE_KEY] = new Map()) +/** Per-workspace flag: has this session ever held a live SSE subscription for it? */ +const SSE_SUBSCRIBED_KEY = '__mcp_sse_subscribed' as const +const sseEverSubscribed: Set = + ((globalThis as Record)[SSE_SUBSCRIBED_KEY] as Set) ?? + ((globalThis as Record)[SSE_SUBSCRIBED_KEY] = new Set()) + /** Subscribes to `tools_changed` SSE events and invalidates the affected query keys. */ export function useMcpToolsEvents(workspaceId: string) { const queryClient = useQueryClient() @@ -563,7 +580,23 @@ export function useMcpToolsEvents(workspaceId: string) { invalidate(serverId) }) + // EventSource fires `onopen` on the initial connect and on every auto-reconnect. Re-sync + // the workspace whenever we could have missed a `tools_changed` event: on any reconnect, + // on the first open of a RE-subscription (leaving the tab tears the connection down), and + // on a first open that only succeeded after an earlier connection error (the initial tools + // query may have failed during that gap and won't retry itself). Skip only a clean first + // subscription — the queries fetch fresh on their own initial mount. + const isResubscribe = sseEverSubscribed.has(workspaceId) + sseEverSubscribed.add(workspaceId) + let opened = false + let erroredBeforeOpen = false + source.onopen = () => { + if (opened || isResubscribe || erroredBeforeOpen) invalidate() + opened = true + } + source.onerror = () => { + if (!opened) erroredBeforeOpen = true logger.warn(`SSE connection error for workspace ${workspaceId}`) } diff --git a/apps/sim/lib/mcp/domain-check.ts b/apps/sim/lib/mcp/domain-check.ts index e04b6f04cc3..1c4c40e3ac1 100644 --- a/apps/sim/lib/mcp/domain-check.ts +++ b/apps/sim/lib/mcp/domain-check.ts @@ -139,14 +139,17 @@ function isLocalhostHostname(hostname: string): boolean { * URLs with env var references in the hostname are skipped — they will be * validated after resolution at execution time. * - * Returns the IP address to pin subsequent connections to (the resolved IP for - * hostnames, or the literal itself for public IP-literal URLs) so the caller can - * prevent DNS-rebinding TOCTOU attacks and stop redirects from escaping to - * internal hosts. Pinning matters for IP literals too: without it the transport - * uses the default fetch, which follows an attacker-controlled 3xx redirect to a - * private/metadata address. Returns null only when pinning is unnecessary or - * impossible: no URL, allowlist-only mode, env-var hostnames (validated later), - * and localhost on self-hosted (no rebinding risk against a fixed loopback). + * Returns the resolved IP (or the literal itself for IP-literal URLs) as a + * non-null **policy signal**: the SSRF guard is active for this server. A public + * resolution selects the validate-at-connect guarded fetch — DNS-rebinding TOCTOU + * and redirect escapes are prevented by re-validating every socket connect and + * following redirects under per-hop validation (see `createSsrfGuardedMcpFetch` / + * `followRedirectsGuarded`), NOT by pinning to this address. The value is literally + * pinned only for the self-hosted private/loopback carve-out (a policy-permitted + * DNS alias the guarded lookup would otherwise filter). Returns null when the guard + * is unnecessary or impossible: no URL, allowlist-only mode, env-var hostnames + * (validated later), and localhost on self-hosted (no rebinding risk against a + * fixed loopback). * * @throws McpSsrfError if the URL resolves to a blocked IP address */