⚒️ fix: MCP Initialization Flows (#8734)

* 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)
This commit is contained in:
Dustin Healy 2025-07-29 11:54:07 -07:00 committed by GitHub
parent 6671fcb714
commit 6fd3b569ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 73 additions and 23 deletions

View file

@ -303,7 +303,7 @@ router.post('/oauth/cancel/:serverName', requireJwtAuth, async (req, res) => {
} }
// Cancel the flow by marking it as failed // 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}`); 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({ res.json({
success: userConnection && !oauthRequired, success: (userConnection && !oauthRequired) || (oauthRequired && oauthUrl),
message: getResponseMessage(), message: getResponseMessage(),
serverName, serverName,
oauthRequired, oauthRequired,

View file

@ -287,14 +287,26 @@ async function checkOAuthFlowStatus(userId, serverName) {
const flowTTL = flowState.ttl || 180000; // Default 3 minutes const flowTTL = flowState.ttl || 180000; // Default 3 minutes
if (flowState.status === 'FAILED' || flowAge > flowTTL) { if (flowState.status === 'FAILED' || flowAge > flowTTL) {
logger.debug(`[MCP Connection Status] Found failed OAuth flow for ${serverName}`, { const wasCancelled = flowState.error && flowState.error.includes('cancelled');
flowId,
status: flowState.status, if (wasCancelled) {
flowAge, logger.debug(`[MCP Connection Status] Found cancelled OAuth flow for ${serverName}`, {
flowTTL, flowId,
timedOut: flowAge > flowTTL, status: flowState.status,
}); error: flowState.error,
return { hasActiveFlow: false, hasFailedFlow: true }; });
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') { if (flowState.status === 'PENDING') {

View file

@ -13,6 +13,7 @@ function MCPSelect() {
batchToggleServers, batchToggleServers,
getServerStatusIconProps, getServerStatusIconProps,
getConfigDialogProps, getConfigDialogProps,
isInitializing,
localize, localize,
} = useMCPServerManager(); } = useMCPServerManager();
@ -32,14 +33,18 @@ function MCPSelect() {
const renderItemContent = useCallback( const renderItemContent = useCallback(
(serverName: string, defaultContent: React.ReactNode) => { (serverName: string, defaultContent: React.ReactNode) => {
const statusIconProps = getServerStatusIconProps(serverName); const statusIconProps = getServerStatusIconProps(serverName);
const isServerInitializing = isInitializing(serverName);
// Common wrapper for the main content (check mark + text) // Common wrapper for the main content (check mark + text)
// Ensures Check & Text are adjacent and the group takes available space. // Ensures Check & Text are adjacent and the group takes available space.
const mainContentWrapper = ( const mainContentWrapper = (
<button <button
type="button" type="button"
className="flex flex-grow items-center rounded bg-transparent p-0 text-left transition-colors focus:outline-none" className={`flex flex-grow items-center rounded bg-transparent p-0 text-left transition-colors focus:outline-none ${
isServerInitializing ? 'opacity-50' : ''
}`}
tabIndex={0} tabIndex={0}
disabled={isServerInitializing}
> >
{defaultContent} {defaultContent}
</button> </button>
@ -58,7 +63,7 @@ function MCPSelect() {
return mainContentWrapper; return mainContentWrapper;
}, },
[getServerStatusIconProps], [getServerStatusIconProps, isInitializing],
); );
// Don't render if no servers are selected and not pinned // Don't render if no servers are selected and not pinned

View file

@ -22,6 +22,7 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
toggleServerSelection, toggleServerSelection,
getServerStatusIconProps, getServerStatusIconProps,
getConfigDialogProps, getConfigDialogProps,
isInitializing,
} = useMCPServerManager(); } = useMCPServerManager();
const menuStore = Ariakit.useMenuStore({ const menuStore = Ariakit.useMenuStore({
@ -86,6 +87,7 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
{configuredServers.map((serverName) => { {configuredServers.map((serverName) => {
const statusIconProps = getServerStatusIconProps(serverName); const statusIconProps = getServerStatusIconProps(serverName);
const isSelected = mcpValues?.includes(serverName) ?? false; const isSelected = mcpValues?.includes(serverName) ?? false;
const isServerInitializing = isInitializing(serverName);
const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />; const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />;
@ -96,12 +98,15 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
event.preventDefault(); event.preventDefault();
toggleServerSelection(serverName); toggleServerSelection(serverName);
}} }}
disabled={isServerInitializing}
className={cn( className={cn(
'flex items-center gap-2 rounded-lg px-2 py-1.5 text-text-primary hover:cursor-pointer', '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', 'scroll-m-1 outline-none transition-colors',
'hover:bg-black/[0.075] dark:hover:bg-white/10', 'hover:bg-black/[0.075] dark:hover:bg-white/10',
'data-[active-item]:bg-black/[0.075] dark:data-[active-item]: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', 'w-full min-w-0 justify-between text-sm',
isServerInitializing &&
'opacity-50 hover:bg-transparent dark:hover:bg-transparent',
)} )}
> >
<div className="flex flex-grow items-center gap-2"> <div className="flex flex-grow items-center gap-2">

View file

@ -34,7 +34,7 @@ export default function ServerInitializationSection({
const serverOAuthUrl = getOAuthUrl(serverName); const serverOAuthUrl = getOAuthUrl(serverName);
const handleInitializeClick = useCallback(() => { const handleInitializeClick = useCallback(() => {
initializeServer(serverName); initializeServer(serverName, false);
}, [initializeServer, serverName]); }, [initializeServer, serverName]);
const handleCancelClick = useCallback(() => { const handleCancelClick = useCallback(() => {

View file

@ -171,6 +171,7 @@ export function useMCPServerManager() {
message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }), message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }),
status: 'error', status: 'error',
}); });
clearInterval(pollInterval);
cleanupServerState(serverName); cleanupServerState(serverName);
return; return;
} }
@ -180,10 +181,15 @@ export function useMCPServerManager() {
message: localize('com_ui_mcp_init_failed'), message: localize('com_ui_mcp_init_failed'),
status: 'error', status: 'error',
}); });
clearInterval(pollInterval);
cleanupServerState(serverName); cleanupServerState(serverName);
return;
} }
} catch (error) { } catch (error) {
console.error(`[MCP Manager] Error polling server ${serverName}:`, error); console.error(`[MCP Manager] Error polling server ${serverName}:`, error);
clearInterval(pollInterval);
cleanupServerState(serverName);
return;
} }
}, 3500); }, 3500);
@ -201,7 +207,7 @@ export function useMCPServerManager() {
); );
const initializeServer = useCallback( const initializeServer = useCallback(
async (serverName: string) => { async (serverName: string, autoOpenOAuth: boolean = true) => {
updateServerState(serverName, { isInitializing: true }); updateServerState(serverName, { isInitializing: true });
try { try {
@ -216,7 +222,9 @@ export function useMCPServerManager() {
isInitializing: true, isInitializing: true,
}); });
window.open(response.oauthUrl, '_blank', 'noopener,noreferrer'); if (autoOpenOAuth) {
window.open(response.oauthUrl, '_blank', 'noopener,noreferrer');
}
startServerPolling(serverName); startServerPolling(serverName);
} else { } else {
@ -265,13 +273,25 @@ export function useMCPServerManager() {
const cancelOAuthFlow = useCallback( const cancelOAuthFlow = useCallback(
(serverName: string) => { (serverName: string) => {
queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]); // Call backend cancellation first, then clean up frontend state on success
cleanupServerState(serverName); cancelOAuthMutation.mutate(serverName, {
cancelOAuthMutation.mutate(serverName); onSuccess: () => {
// Only clean up frontend state after backend confirms cancellation
cleanupServerState(serverName);
queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]);
showToast({ showToast({
message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }), message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }),
status: 'warning', 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], [queryClient, cleanupServerState, showToast, localize, cancelOAuthMutation],
@ -309,6 +329,10 @@ export function useMCPServerManager() {
const disconnectedServers: string[] = []; const disconnectedServers: string[] = [];
serverNames.forEach((serverName) => { serverNames.forEach((serverName) => {
if (isInitializing(serverName)) {
return;
}
const serverStatus = connectionStatus[serverName]; const serverStatus = connectionStatus[serverName];
if (serverStatus?.connectionState === 'connected') { if (serverStatus?.connectionState === 'connected') {
connectedServers.push(serverName); connectedServers.push(serverName);
@ -323,11 +347,15 @@ export function useMCPServerManager() {
initializeServer(serverName); initializeServer(serverName);
}); });
}, },
[connectionStatus, setMCPValues, initializeServer], [connectionStatus, setMCPValues, initializeServer, isInitializing],
); );
const toggleServerSelection = useCallback( const toggleServerSelection = useCallback(
(serverName: string) => { (serverName: string) => {
if (isInitializing(serverName)) {
return;
}
const currentValues = mcpValues ?? []; const currentValues = mcpValues ?? [];
const isCurrentlySelected = currentValues.includes(serverName); const isCurrentlySelected = currentValues.includes(serverName);
@ -343,7 +371,7 @@ export function useMCPServerManager() {
} }
} }
}, },
[mcpValues, setMCPValues, connectionStatus, initializeServer], [mcpValues, setMCPValues, connectionStatus, initializeServer, isInitializing],
); );
const handleConfigSave = useCallback( const handleConfigSave = useCallback(