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 92f3fc7f60a..9708b1b1ff9 100644 --- a/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/mcp.tsx +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mcp/mcp.tsx @@ -68,10 +68,12 @@ interface ServerListItemProps { server: McpServer tools: McpTool[] isDeleting: boolean + isConnecting: boolean isLoadingTools?: boolean isRefreshing?: boolean onRemove: () => void onViewDetails: () => void + onAuthorize: () => void } function ServerListItem({ @@ -79,10 +81,12 @@ function ServerListItem({ server, tools, isDeleting, + isConnecting, isLoadingTools = false, isRefreshing = false, onRemove, onViewDetails, + onAuthorize, }: ServerListItemProps) { const transportLabel = formatTransportLabel(server.transport || 'http') const toolsLabel = getServerToolsLabel( @@ -121,6 +125,14 @@ function ServerListItem({ label='Server actions' actions={[ { label: 'Details', onSelect: onViewDetails }, + ...(canManage && server.authType === 'oauth' && server.connectionStatus !== 'connected' + ? [ + { + label: isConnecting ? 'Reopen authorization' : 'Authorize', + onSelect: onAuthorize, + }, + ] + : []), ...(canManage ? [ { @@ -450,12 +462,13 @@ export function MCP() {
{ await startOauthForServer(server.id) }} > - {connectingOauthServers.has(server.id) ? 'Connecting…' : 'Connect with OAuth'} + {connectingOauthServers.has(server.id) + ? 'Reopen authorization window' + : 'Connect with OAuth'}
@@ -660,6 +673,7 @@ export function MCP() { server={server} tools={tools} isDeleting={deletingServers.has(server.id)} + isConnecting={connectingOauthServers.has(server.id)} isLoadingTools={isLoadingTools} isRefreshing={ refreshServerMutation.isPending && @@ -667,6 +681,7 @@ export function MCP() { } onRemove={() => handleRemoveServer(server.id)} onViewDetails={() => handleViewDetails(server.id)} + onAuthorize={() => startOauthForServer(server.id)} /> ) })} diff --git a/apps/sim/hooks/mcp/use-mcp-oauth-popup.test.tsx b/apps/sim/hooks/mcp/use-mcp-oauth-popup.test.tsx new file mode 100644 index 00000000000..234a71b057d --- /dev/null +++ b/apps/sim/hooks/mcp/use-mcp-oauth-popup.test.tsx @@ -0,0 +1,191 @@ +/** + * @vitest-environment jsdom + */ +import { act, type ReactNode } from 'react' +import { sleep } from '@sim/utils/helpers' +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { createRoot, type Root } from 'react-dom/client' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockStartOauth } = vi.hoisted(() => ({ mockStartOauth: vi.fn() })) + +vi.mock('@sim/emcn', () => ({ + toast: { success: vi.fn(), error: vi.fn() }, +})) + +vi.mock('@/hooks/queries/mcp', () => ({ + useStartMcpOauth: () => ({ mutateAsync: mockStartOauth }), + mcpKeys: { + serversList: (workspaceId: string) => ['mcp', 'servers', workspaceId], + serverToolsList: (workspaceId: string, serverId: string) => [ + 'mcp', + 'server-tools', + workspaceId, + serverId, + ], + storedToolsList: (workspaceId: string) => ['mcp', 'stored-tools', workspaceId], + }, +})) + +import { useMcpOauthPopup } from '@/hooks/mcp/use-mcp-oauth-popup' + +/** + * Minimal dependency-free hook harness (the repo has no `@testing-library/react`). + * Mounts the hook in a real React 19 root under jsdom, wrapped in a real + * `QueryClientProvider`, so query/mutation lifecycles run exactly as in the app. + */ +function renderHookWithClient(useHook: () => T): { + result: () => T + queryClient: QueryClient + unmount: () => void +} { + ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }) + const container = document.createElement('div') + const root: Root = createRoot(container) + let latest: T + + function Probe() { + latest = useHook() + return null + } + + function Wrapper({ children }: { children: ReactNode }) { + return {children} + } + + act(() => { + root.render( + + + + ) + }) + + return { result: () => latest, queryClient, unmount: () => act(() => root.unmount()) } +} + +async function flush() { + await act(async () => { + for (let i = 0; i < 5; i++) { + await Promise.resolve() + await sleep(0) + } + }) +} + +describe('useMcpOauthPopup', () => { + beforeEach(() => { + vi.clearAllMocks() + // jsdom has no BroadcastChannel; the hook opens one on mount. + class FakeBroadcastChannel { + onmessage: ((event: MessageEvent) => void) | null = null + constructor(public name: string) {} + postMessage(): void {} + close(): void {} + } + ;(globalThis as unknown as { BroadcastChannel: unknown }).BroadcastChannel = + FakeBroadcastChannel + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + it('ignores a concurrent second start for the same server (no double popup)', async () => { + let resolveStart: (value: unknown) => void = () => {} + mockStartOauth.mockImplementation( + () => + new Promise((res) => { + resolveStart = res + }) + ) + + const hook = renderHookWithClient(() => useMcpOauthPopup({ workspaceId: 'w1' })) + await flush() + + // Two clicks before /oauth/start resolves — the guard must collapse them to one request. + await act(async () => { + void hook.result().startOauthForServer('s1') + void hook.result().startOauthForServer('s1') + }) + expect(mockStartOauth).toHaveBeenCalledTimes(1) + + // Settle the first flow so the guard clears. + await act(async () => { + resolveStart({ status: 'redirect', popup: { closed: false }, state: 'state-1' }) + }) + await flush() + + hook.unmount() + }) + + it('allows a fresh start after the previous one settles (reopen after abandon)', async () => { + mockStartOauth.mockResolvedValue({ status: 'redirect', popup: { closed: false }, state: 'st' }) + + const hook = renderHookWithClient(() => useMcpOauthPopup({ workspaceId: 'w1' })) + await flush() + + await act(async () => { + await hook.result().startOauthForServer('s1') + }) + await flush() + + await act(async () => { + await hook.result().startOauthForServer('s1') + }) + await flush() + + // Both distinct clicks reached the mutation — the guard only blocks concurrent re-entry. + expect(mockStartOauth).toHaveBeenCalledTimes(2) + + hook.unmount() + }) + + it('invalidates server queries when start reports already_authorized', async () => { + mockStartOauth.mockResolvedValue({ status: 'already_authorized' }) + + const hook = renderHookWithClient(() => useMcpOauthPopup({ workspaceId: 'w1' })) + await flush() + const invalidateSpy = vi.spyOn(hook.queryClient, 'invalidateQueries') + + await act(async () => { + await hook.result().startOauthForServer('s1') + }) + await flush() + + // The server is already connected — the UI must still refresh so it reflects that. + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['mcp', 'servers', 'w1'] }) + + hook.unmount() + }) + + it('keeps the row connecting when a reopen fails but a prior flow is still live', async () => { + mockStartOauth.mockResolvedValueOnce({ + status: 'redirect', + popup: { closed: false }, + state: 'st-a', + }) + const hook = renderHookWithClient(() => useMcpOauthPopup({ workspaceId: 'w1' })) + await flush() + + await act(async () => { + await hook.result().startOauthForServer('s1') + }) + await flush() + expect(hook.result().connectingServers.has('s1')).toBe(true) + + // Reopen fails (e.g. popup blocked). The still-live prior flow must keep the row connecting + // rather than the failed attempt clearing it (deterministic reference counting). + mockStartOauth.mockRejectedValueOnce(new Error('Popup blocked')) + await act(async () => { + await hook.result().startOauthForServer('s1') + }) + await flush() + expect(hook.result().connectingServers.has('s1')).toBe(true) + + hook.unmount() + }) +}) diff --git a/apps/sim/hooks/mcp/use-mcp-oauth-popup.ts b/apps/sim/hooks/mcp/use-mcp-oauth-popup.ts index a7588ba7bac..678d9ae7c1b 100644 --- a/apps/sim/hooks/mcp/use-mcp-oauth-popup.ts +++ b/apps/sim/hooks/mcp/use-mcp-oauth-popup.ts @@ -1,6 +1,6 @@ 'use client' -import { useCallback, useEffect, useRef, useState } from 'react' +import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { toast } from '@sim/emcn' import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' @@ -48,136 +48,157 @@ export function useMcpOauthPopup({ workspaceId }: UseMcpOauthPopupProps) { const queryClient = useQueryClient() const { mutateAsync: startOauth } = useStartMcpOauth() - const [connectingServers, setConnectingServers] = useState>(() => new Set()) - // OAuth `state` nonce -> { serverId, safety timeout }. The state keys the BroadcastChannel - // correlation: the callback echoes it on every result (even failures that can't resolve a - // serverId), so the tab that started this exact flow matches it while other same-origin tabs - // ignore it. Cleared only when the flow completes or times out, never by popup.closed polling - // — COOP can make popup.closed misreport, and clearing early would drop a genuine completion. + // Per-server count of live authorization attempts; a row shows "Connecting…" / "Reopen + // authorization" while its count > 0. Reference counting (not a boolean set) keeps the label + // deterministic across concurrent attempts: a reopen increments before the superseded flow + // decrements, so the count never dips to 0 mid-reopen (no flicker), and every attempt clears + // exactly once (never stuck). + const [connectingCounts, setConnectingCounts] = useState>(() => new Map()) + // OAuth `state` nonce -> { serverId, safety timeout }. `state` keys the BroadcastChannel + // correlation: the callback echoes it on every result (even failures that resolve no serverId), + // so the tab that started this exact flow matches it while other same-origin tabs — and + // unrelated flows in this tab — ignore the broadcast. const pendingFlowsRef = useRef>(new Map()) - // serverId -> popup.closed poll. Best-effort fast "Connecting…" clear when the user - // abandons the popup; never used to correlate a result. - const popupPollsRef = useRef>(new Map()) - - const stopConnecting = useCallback((serverId: string) => { - setConnectingServers((prev) => { - if (!prev.has(serverId)) return prev - const next = new Set(prev) - next.delete(serverId) + // serverIds with an in-flight `/oauth/start` request — guards a fast double-click from opening + // two popups. Cleared once the request settles, so a later click (to reopen an abandoned + // popup) still starts a fresh flow. + const startingRef = useRef | null>(null) + + const incConnecting = useCallback((serverId: string) => { + setConnectingCounts((prev) => { + const next = new Map(prev) + next.set(serverId, (next.get(serverId) ?? 0) + 1) return next }) }, []) - const stopPopupPoll = useCallback((serverId: string) => { - const poll = popupPollsRef.current.get(serverId) - if (poll !== undefined) { - window.clearInterval(poll) - popupPollsRef.current.delete(serverId) - } + const decConnecting = useCallback((serverId: string) => { + setConnectingCounts((prev) => { + const current = prev.get(serverId) + if (current === undefined) return prev + const next = new Map(prev) + if (current <= 1) next.delete(serverId) + else next.set(serverId, current - 1) + return next + }) }, []) - /** End a flow entirely (by its state nonce): stop the spinner, safety timeout, and popup poll. */ + const invalidateServer = useCallback( + (serverId: string) => { + queryClient.invalidateQueries({ queryKey: mcpKeys.serversList(workspaceId) }) + queryClient.invalidateQueries({ queryKey: mcpKeys.serverToolsList(workspaceId, serverId) }) + queryClient.invalidateQueries({ queryKey: mcpKeys.storedToolsList(workspaceId) }) + }, + [queryClient, workspaceId] + ) + + /** End one flow by its `state` nonce, decrementing its server's connecting count exactly once. */ const settleFlow = useCallback( (state: string) => { const flow = pendingFlowsRef.current.get(state) if (!flow) return window.clearTimeout(flow.timeout) pendingFlowsRef.current.delete(state) - stopPopupPoll(flow.serverId) - stopConnecting(flow.serverId) + decConnecting(flow.serverId) + }, + [decConnecting] + ) + + /** Retire every prior flow for a server, superseded by a fresh attempt (each decrements once). */ + const retireFlows = useCallback( + (serverId: string) => { + const states: string[] = [] + for (const [state, flow] of pendingFlowsRef.current) { + if (flow.serverId === serverId) states.push(state) + } + for (const state of states) settleFlow(state) }, - [stopConnecting, stopPopupPoll] + [settleFlow] ) useEffect(() => { const pending = pendingFlowsRef.current - const polls = popupPollsRef.current return () => { for (const { timeout } of pending.values()) window.clearTimeout(timeout) - for (const p of polls.values()) window.clearInterval(p) pending.clear() - polls.clear() } }, []) useEffect(() => { - // The callback signals over a same-origin BroadcastChannel (see the OAuth callback - // route): a provider whose authorize page sets COOP `same-origin` severs - // `window.opener`, so a popup `postMessage` can be lost and leave the row stuck on - // "Connecting…". A BroadcastChannel is origin-scoped, so it needs no origin check. + // The callback signals over a same-origin BroadcastChannel (see the OAuth callback route): a + // provider whose authorize page sets COOP `same-origin` severs `window.opener`, so a popup + // `postMessage` can be lost and leave the row stuck on "Connecting…". A BroadcastChannel is + // origin-scoped, so it needs no origin check. const channel = new BroadcastChannel('mcp-oauth') channel.onmessage = (event) => { const data = event.data as Partial | null if (data?.type !== 'mcp-oauth') return - // A BroadcastChannel reaches every same-origin tab, so react only to a result for a flow - // THIS tab started, matched on the OAuth `state` nonce. Every result (success or failure) - // carries it, so unrelated tabs — and unrelated flows in this tab — ignore the broadcast. if (!data.state) return const flow = pendingFlowsRef.current.get(data.state) if (!flow) return const { serverId } = flow settleFlow(data.state) if (data.ok) { - queryClient.invalidateQueries({ queryKey: mcpKeys.serversList(workspaceId) }) - queryClient.invalidateQueries({ - queryKey: mcpKeys.serverToolsList(workspaceId, serverId), - }) - queryClient.invalidateQueries({ queryKey: mcpKeys.storedToolsList(workspaceId) }) + invalidateServer(serverId) toast.success('Server authorized') } else { toast.error(reasonToMessage(data.reason)) } } return () => channel.close() - }, [queryClient, workspaceId, settleFlow]) + }, [settleFlow, invalidateServer]) const startOauthForServer = useCallback( async (serverId: string) => { - setConnectingServers((prev) => new Set(prev).add(serverId)) + const starting = (startingRef.current ??= new Set()) + if (starting.has(serverId)) return + starting.add(serverId) + incConnecting(serverId) // this attempt begins try { const result = await startOauth({ serverId, workspaceId }) + // The replacement start succeeded (already-authorized, or a fresh popup opened), so retire + // any prior attempt for this server now — its result is moot and the server-side `state` + // it depended on has been overwritten. A *failed* start (below) leaves prior flows intact. + retireFlows(serverId) if (result.status === 'already_authorized') { - stopConnecting(serverId) + invalidateServer(serverId) + decConnecting(serverId) // this attempt ends return } - const { popup, state } = result - // Drop any prior in-flight flow for this server (e.g. an abandoned attempt now being - // retried) so its stale safety timeout can't later clear this new flow's state. - for (const [prevState, flow] of pendingFlowsRef.current) { - if (flow.serverId === serverId) { - window.clearTimeout(flow.timeout) - pendingFlowsRef.current.delete(prevState) - } - } - // Track this in-flight flow keyed by its `state` nonce for the BroadcastChannel gate, - // bounded by a safety timeout in case no result ever arrives (popup abandoned, or a - // callback failure the client can't otherwise clear). + // Track the in-flight flow by its `state` nonce for the BroadcastChannel gate, bounded by + // a safety timeout in case no result ever arrives (popup abandoned, or a callback the + // client can't otherwise observe under COOP). + const { state } = result pendingFlowsRef.current.set(state, { serverId, timeout: window.setTimeout(() => settleFlow(state), OAUTH_FLOW_TIMEOUT_MS), }) - // Best-effort: clear "Connecting…" quickly when the user closes the popup without - // finishing. popup.closed can misreport under COOP, so this only stops the spinner — - // it never touches `pendingFlowsRef`, so it can't drop a real result. - stopPopupPoll(serverId) - popupPollsRef.current.set( - serverId, - window.setInterval(() => { - if (popup.closed) { - stopPopupPoll(serverId) - stopConnecting(serverId) - } - }, 500) - ) } catch (e) { - stopPopupPoll(serverId) - stopConnecting(serverId) + decConnecting(serverId) // this attempt ends; any prior flow keeps its own count logger.error('Failed to start MCP OAuth', e) toast.error(toError(e).message || 'Failed to start authorization') + } finally { + starting.delete(serverId) } }, - [startOauth, workspaceId, settleFlow, stopConnecting, stopPopupPoll] + [ + startOauth, + workspaceId, + settleFlow, + retireFlows, + incConnecting, + decConnecting, + invalidateServer, + ] ) + const connectingServers = useMemo(() => { + const set = new Set() + for (const [serverId, count] of connectingCounts) { + if (count > 0) set.add(serverId) + } + return set + }, [connectingCounts]) + return { connectingServers, startOauthForServer } }