⚒️ 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
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,

View file

@ -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') {

View file

@ -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 = (
<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}
disabled={isServerInitializing}
>
{defaultContent}
</button>
@ -58,7 +63,7 @@ function MCPSelect() {
return mainContentWrapper;
},
[getServerStatusIconProps],
[getServerStatusIconProps, isInitializing],
);
// Don't render if no servers are selected and not pinned

View file

@ -22,6 +22,7 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
toggleServerSelection,
getServerStatusIconProps,
getConfigDialogProps,
isInitializing,
} = useMCPServerManager();
const menuStore = Ariakit.useMenuStore({
@ -86,6 +87,7 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
{configuredServers.map((serverName) => {
const statusIconProps = getServerStatusIconProps(serverName);
const isSelected = mcpValues?.includes(serverName) ?? false;
const isServerInitializing = isInitializing(serverName);
const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />;
@ -96,12 +98,15 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
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',
)}
>
<div className="flex flex-grow items-center gap-2">

View file

@ -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(() => {

View file

@ -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(