From 6fd3b569accac0986afe9c18c390da081e116169 Mon Sep 17 00:00:00 2001 From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> Date: Tue, 29 Jul 2025 11:54:07 -0700 Subject: [PATCH] =?UTF-8?q?=E2=9A=92=EF=B8=8F=20fix:=20MCP=20Initializatio?= =?UTF-8?q?n=20Flows=20(#8734)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: add OAuth flow back in to success state * feat: disable server clicks during initialization to prevent spam * fix: correct new tab behavior for OAuth between one-click and normal initialization flows * fix: stop polling on error during oauth (was infinite popping toasts because we didn't clear interval) * fix: cleanupServerState should be called after successful cancelOauth, not before * fix: change from completeFlow to failFlow to avoid stale client IDs on OAuth after cancellation * fix: add logic to differentiate between cancelled and failed flows when checking status for indicators (so error triangle indicator doesn't show up on cancellaiton) --- api/server/routes/mcp.js | 4 +- api/server/services/MCP.js | 28 +++++++---- .../src/components/Chat/Input/MCPSelect.tsx | 9 +++- .../src/components/Chat/Input/MCPSubMenu.tsx | 5 ++ .../MCP/ServerInitializationSection.tsx | 2 +- client/src/hooks/MCP/useMCPServerManager.ts | 48 +++++++++++++++---- 6 files changed, 73 insertions(+), 23 deletions(-) diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index f66d671a8..270c8525a 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -303,7 +303,7 @@ router.post('/oauth/cancel/:serverName', requireJwtAuth, async (req, res) => { } // Cancel the flow by marking it as failed - await flowManager.completeFlow(flowId, 'mcp_oauth', null, 'User cancelled OAuth flow'); + await flowManager.failFlow(flowId, 'mcp_oauth', 'User cancelled OAuth flow'); logger.info(`[MCP OAuth Cancel] Successfully cancelled OAuth flow for ${serverName}`); @@ -463,7 +463,7 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => { }; res.json({ - success: userConnection && !oauthRequired, + success: (userConnection && !oauthRequired) || (oauthRequired && oauthUrl), message: getResponseMessage(), serverName, oauthRequired, diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index f8ec2d04d..147def1bb 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -287,14 +287,26 @@ async function checkOAuthFlowStatus(userId, serverName) { const flowTTL = flowState.ttl || 180000; // Default 3 minutes if (flowState.status === 'FAILED' || flowAge > flowTTL) { - logger.debug(`[MCP Connection Status] Found failed OAuth flow for ${serverName}`, { - flowId, - status: flowState.status, - flowAge, - flowTTL, - timedOut: flowAge > flowTTL, - }); - return { hasActiveFlow: false, hasFailedFlow: true }; + const wasCancelled = flowState.error && flowState.error.includes('cancelled'); + + if (wasCancelled) { + logger.debug(`[MCP Connection Status] Found cancelled OAuth flow for ${serverName}`, { + flowId, + status: flowState.status, + error: flowState.error, + }); + return { hasActiveFlow: false, hasFailedFlow: false }; + } else { + logger.debug(`[MCP Connection Status] Found failed OAuth flow for ${serverName}`, { + flowId, + status: flowState.status, + flowAge, + flowTTL, + timedOut: flowAge > flowTTL, + error: flowState.error, + }); + return { hasActiveFlow: false, hasFailedFlow: true }; + } } if (flowState.status === 'PENDING') { diff --git a/client/src/components/Chat/Input/MCPSelect.tsx b/client/src/components/Chat/Input/MCPSelect.tsx index 8160cdd2f..01c4d72a4 100644 --- a/client/src/components/Chat/Input/MCPSelect.tsx +++ b/client/src/components/Chat/Input/MCPSelect.tsx @@ -13,6 +13,7 @@ function MCPSelect() { batchToggleServers, getServerStatusIconProps, getConfigDialogProps, + isInitializing, localize, } = useMCPServerManager(); @@ -32,14 +33,18 @@ function MCPSelect() { const renderItemContent = useCallback( (serverName: string, defaultContent: React.ReactNode) => { const statusIconProps = getServerStatusIconProps(serverName); + const isServerInitializing = isInitializing(serverName); // Common wrapper for the main content (check mark + text) // Ensures Check & Text are adjacent and the group takes available space. const mainContentWrapper = ( @@ -58,7 +63,7 @@ function MCPSelect() { return mainContentWrapper; }, - [getServerStatusIconProps], + [getServerStatusIconProps, isInitializing], ); // Don't render if no servers are selected and not pinned diff --git a/client/src/components/Chat/Input/MCPSubMenu.tsx b/client/src/components/Chat/Input/MCPSubMenu.tsx index 3628b732a..fc690089d 100644 --- a/client/src/components/Chat/Input/MCPSubMenu.tsx +++ b/client/src/components/Chat/Input/MCPSubMenu.tsx @@ -22,6 +22,7 @@ const MCPSubMenu = React.forwardRef( toggleServerSelection, getServerStatusIconProps, getConfigDialogProps, + isInitializing, } = useMCPServerManager(); const menuStore = Ariakit.useMenuStore({ @@ -86,6 +87,7 @@ const MCPSubMenu = React.forwardRef( {configuredServers.map((serverName) => { const statusIconProps = getServerStatusIconProps(serverName); const isSelected = mcpValues?.includes(serverName) ?? false; + const isServerInitializing = isInitializing(serverName); const statusIcon = statusIconProps && ; @@ -96,12 +98,15 @@ const MCPSubMenu = React.forwardRef( event.preventDefault(); toggleServerSelection(serverName); }} + disabled={isServerInitializing} className={cn( 'flex items-center gap-2 rounded-lg px-2 py-1.5 text-text-primary hover:cursor-pointer', 'scroll-m-1 outline-none transition-colors', 'hover:bg-black/[0.075] dark:hover:bg-white/10', 'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10', 'w-full min-w-0 justify-between text-sm', + isServerInitializing && + 'opacity-50 hover:bg-transparent dark:hover:bg-transparent', )} >
diff --git a/client/src/components/MCP/ServerInitializationSection.tsx b/client/src/components/MCP/ServerInitializationSection.tsx index 2113a9f84..0623ba1a2 100644 --- a/client/src/components/MCP/ServerInitializationSection.tsx +++ b/client/src/components/MCP/ServerInitializationSection.tsx @@ -34,7 +34,7 @@ export default function ServerInitializationSection({ const serverOAuthUrl = getOAuthUrl(serverName); const handleInitializeClick = useCallback(() => { - initializeServer(serverName); + initializeServer(serverName, false); }, [initializeServer, serverName]); const handleCancelClick = useCallback(() => { diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index 71ffa471b..b42854545 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -171,6 +171,7 @@ export function useMCPServerManager() { message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }), status: 'error', }); + clearInterval(pollInterval); cleanupServerState(serverName); return; } @@ -180,10 +181,15 @@ export function useMCPServerManager() { message: localize('com_ui_mcp_init_failed'), status: 'error', }); + clearInterval(pollInterval); cleanupServerState(serverName); + return; } } catch (error) { console.error(`[MCP Manager] Error polling server ${serverName}:`, error); + clearInterval(pollInterval); + cleanupServerState(serverName); + return; } }, 3500); @@ -201,7 +207,7 @@ export function useMCPServerManager() { ); const initializeServer = useCallback( - async (serverName: string) => { + async (serverName: string, autoOpenOAuth: boolean = true) => { updateServerState(serverName, { isInitializing: true }); try { @@ -216,7 +222,9 @@ export function useMCPServerManager() { isInitializing: true, }); - window.open(response.oauthUrl, '_blank', 'noopener,noreferrer'); + if (autoOpenOAuth) { + window.open(response.oauthUrl, '_blank', 'noopener,noreferrer'); + } startServerPolling(serverName); } else { @@ -265,13 +273,25 @@ export function useMCPServerManager() { const cancelOAuthFlow = useCallback( (serverName: string) => { - queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]); - cleanupServerState(serverName); - cancelOAuthMutation.mutate(serverName); + // Call backend cancellation first, then clean up frontend state on success + cancelOAuthMutation.mutate(serverName, { + onSuccess: () => { + // Only clean up frontend state after backend confirms cancellation + cleanupServerState(serverName); + queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]); - showToast({ - message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }), - status: 'warning', + showToast({ + message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }), + status: 'warning', + }); + }, + onError: (error) => { + console.error(`[MCP Manager] Failed to cancel OAuth for ${serverName}:`, error); + showToast({ + message: localize('com_ui_mcp_init_failed', { 0: serverName }), + status: 'error', + }); + }, }); }, [queryClient, cleanupServerState, showToast, localize, cancelOAuthMutation], @@ -309,6 +329,10 @@ export function useMCPServerManager() { const disconnectedServers: string[] = []; serverNames.forEach((serverName) => { + if (isInitializing(serverName)) { + return; + } + const serverStatus = connectionStatus[serverName]; if (serverStatus?.connectionState === 'connected') { connectedServers.push(serverName); @@ -323,11 +347,15 @@ export function useMCPServerManager() { initializeServer(serverName); }); }, - [connectionStatus, setMCPValues, initializeServer], + [connectionStatus, setMCPValues, initializeServer, isInitializing], ); const toggleServerSelection = useCallback( (serverName: string) => { + if (isInitializing(serverName)) { + return; + } + const currentValues = mcpValues ?? []; const isCurrentlySelected = currentValues.includes(serverName); @@ -343,7 +371,7 @@ export function useMCPServerManager() { } } }, - [mcpValues, setMCPValues, connectionStatus, initializeServer], + [mcpValues, setMCPValues, connectionStatus, initializeServer, isInitializing], ); const handleConfigSave = useCallback(