feat(ui): de-jank context menu

There was a lot of convoluted, janky logic related to trying to not mount the context menu's portal until its needed. This was in the library where the component was originally copied from.

I've removed that and resolved the jank, at the cost of there being an extra portal for each instance of the context menu. Don't think this is going to be an issue. If it is, the whole context menu could be refactored to be a singleton.

Author: psychedelicious

invokeai/frontend/web/src/common/components/InvContextMenu/InvContextMenu.tsx

@@ -1,26 +1,15 @@
 /**
- * This is a copy-paste of https://github.com/lukasbach/chakra-ui-contextmenu with a small change.
- *
- * The reactflow background element somehow prevents the chakra `useOutsideClick()` hook from working.
- * With a menu open, clicking on the reactflow background element doesn't close the menu.
- *
- * Reactflow does provide an `onPaneClick` to handle clicks on the background element, but it is not
- * straightforward to programatically close the menu.
- *
- * As a (hopefully temporary) workaround, we will use a dirty hack:
- * - create `globalContextMenuCloseTrigger: number` in `ui` slice
- * - increment it in `onPaneClick` (and wherever else we want to close the menu)
- * - `useEffect()` to close the menu when `globalContextMenuCloseTrigger` changes
+ * Adapted from https://github.com/lukasbach/chakra-ui-contextmenu
  */
 import type {
   ChakraProps,
   MenuButtonProps,
   MenuProps,
   PortalProps,
 } from '@chakra-ui/react';
-import { Portal, useEventListener } from '@chakra-ui/react';
+import { Portal, useDisclosure, useEventListener } from '@chakra-ui/react';
 import { InvMenu, InvMenuButton } from 'common/components/InvMenu/wrapper';
-import { useGlobalMenuCloseTrigger } from 'common/hooks/useGlobalMenuCloseTrigger';
+import { useGlobalMenuClose } from 'common/hooks/useGlobalMenuClose';
 import { typedMemo } from 'common/util/typedMemo';
 import { useCallback, useEffect, useRef, useState } from 'react';
 
@@ -34,94 +23,89 @@ export interface InvContextMenuProps<T extends HTMLElement = HTMLDivElement> {
 
 export const InvContextMenu = typedMemo(
   <T extends HTMLElement = HTMLElement>(props: InvContextMenuProps<T>) => {
-    const [isOpen, setIsOpen] = useState(false);
-    const [isRendered, setIsRendered] = useState(false);
-    const [isDeferredOpen, setIsDeferredOpen] = useState(false);
-    const [position, setPosition] = useState<[number, number]>([0, 0]);
+    const { isOpen, onOpen, onClose } = useDisclosure();
+    const [position, setPosition] = useState([-1, -1]);
     const targetRef = useRef<T>(null);
+    const lastPositionRef = useRef([-1, -1]);
+    const timeoutRef = useRef(0);
 
-    useEffect(() => {
-      if (isOpen) {
-        setTimeout(() => {
-          setIsRendered(true);
-          setTimeout(() => {
-            setIsDeferredOpen(true);
-          });
-        });
-      } else {
-        setIsDeferredOpen(false);
-        const timeout = setTimeout(() => {
-          setIsRendered(isOpen);
-        }, 1000);
-        return () => clearTimeout(timeout);
-      }
-    }, [isOpen]);
+    useGlobalMenuClose(onClose);
 
-    const onClose = useCallback(() => {
-      setIsOpen(false);
-      setIsDeferredOpen(false);
-      setIsRendered(false);
-    }, []);
-
-    // This is the change from the original chakra-ui-contextmenu
-    // Close all menus when the globalContextMenuCloseTrigger changes
-    useGlobalMenuCloseTrigger(onClose);
+    const onContextMenu = useCallback(
+      (e: MouseEvent) => {
+        if (e.shiftKey) {
+          onClose();
+          return;
+        }
+        if (
+          targetRef.current?.contains(e.target as HTMLElement) ||
+          e.target === targetRef.current
+        ) {
+          // clear pending delayed open
+          window.clearTimeout(timeoutRef.current);
+          e.preventDefault();
+          if (
+            lastPositionRef.current[0] !== e.pageX ||
+            lastPositionRef.current[1] !== e.pageY
+          ) {
+            // if the mouse moved, we need to close, wait for animation and reopen the menu at the new position
+            onClose();
+            timeoutRef.current = window.setTimeout(() => {
+              onOpen();
+              setPosition([e.pageX, e.pageY]);
+            }, 100);
+          } else {
+            // else we can just open the menu at the current position
+            onOpen();
+            setPosition([e.pageX, e.pageY]);
+          }
+        }
+        lastPositionRef.current = [e.pageX, e.pageY];
+      },
+      [onClose, onOpen]
+    );
 
-    useEventListener('contextmenu', (e) => {
-      if (
-        targetRef.current?.contains(e.target as HTMLElement) ||
-        e.target === targetRef.current
-      ) {
-        e.preventDefault();
-        setIsOpen(true);
-        setPosition([e.pageX, e.pageY]);
-      } else {
-        setIsOpen(false);
-      }
-    });
+    useEffect(
+      () => () => {
+        window.clearTimeout(timeoutRef.current);
+      },
+      []
+    );
 
-    const onCloseHandler = useCallback(() => {
-      props.menuProps?.onClose?.();
-      setIsOpen(false);
-    }, [props.menuProps]);
+    useEventListener('contextmenu', onContextMenu);
 
     return (
       <>
         {props.children(targetRef)}
-        {isRendered && (
-          <Portal {...props.portalProps}>
-            <InvMenu
-              isLazy
-              isOpen={isDeferredOpen}
-              gutter={0}
-              onClose={onCloseHandler}
-              placement="auto-end"
-              {...props.menuProps}
-            >
-              <InvMenuButton
-                aria-hidden={true}
-                w={1}
-                h={1}
-                position="absolute"
-                left={position[0]}
-                top={position[1]}
-                cursor="default"
-                bg="transparent"
-                size="sm"
-                _hover={_hover}
-                {...props.menuButtonProps}
-              />
-              {props.renderMenu()}
-            </InvMenu>
-          </Portal>
-        )}
+        <Portal {...props.portalProps}>
+          <InvMenu
+            isLazy
+            isOpen={isOpen}
+            gutter={0}
+            placement="auto-end"
+            onClose={onClose}
+            {...props.menuProps}
+          >
+            <InvMenuButton
+              aria-hidden={true}
+              w={1}
+              h={1}
+              position="absolute"
+              left={position[0]}
+              top={position[1]}
+              cursor="default"
+              bg="transparent"
+              size="sm"
+              _hover={_hover}
+              pointerEvents="none"
+              {...props.menuButtonProps}
+            />
+            {props.renderMenu()}
+          </InvMenu>
+        </Portal>
       </>
     );
   }
 );
 
 const _hover: ChakraProps['_hover'] = { bg: 'transparent' };
-
-Object.assign(InvContextMenu, {
-  displayName: 'InvContextMenu',
-});

invokeai/frontend/web/src/common/components/InvMenu/InvMenuList.tsx

@@ -3,6 +3,7 @@ import {
   MenuList as ChakraMenuList,
   Portal,
 } from '@chakra-ui/react';
+import { skipMouseEvent } from 'common/util/skipMouseEvent';
 import { memo } from 'react';
 
 import { menuListMotionProps } from './constants';
@@ -16,6 +17,7 @@ export const InvMenuList = memo(
           <ChakraMenuList
             ref={ref}
             motionProps={menuListMotionProps}
+            onContextMenu={skipMouseEvent}
             {...props}
           />
         </Portal>

invokeai/frontend/web/src/common/hooks/useGlobalMenuCloseTrigger.ts renamed to invokeai/frontend/web/src/common/hooks/useGlobalMenuClose.ts

@@ -15,7 +15,7 @@ const $onCloseCallbacks = atom<CB[]>([]);
  * This hook provides a way to close all menus by calling `onCloseGlobal()`. Menus that want to be closed
  * in this way should register themselves by passing a callback to `useGlobalMenuCloseTrigger()`.
  */
-export const useGlobalMenuCloseTrigger = (onClose?: CB) => {
+export const useGlobalMenuClose = (onClose?: CB) => {
   useEffect(() => {
     if (!onClose) {
       return;

invokeai/frontend/web/src/common/util/skipMouseEvent.ts

@@ -0,0 +1,8 @@
+import type { MouseEvent } from 'react';
+
+/**
+ * Prevents the default behavior of the event.
+ */
+export const skipMouseEvent = (e: MouseEvent) => {
+  e.preventDefault();
+};

invokeai/frontend/web/src/features/gallery/components/Boards/BoardContextMenu.tsx

@@ -12,7 +12,6 @@ import {
 import type { BoardId } from 'features/gallery/store/types';
 import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
 import { addToast } from 'features/system/store/systemSlice';
-import type { MouseEvent } from 'react';
 import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { FaDownload, FaPlus } from 'react-icons/fa';
@@ -90,13 +89,9 @@ const BoardContextMenu = ({
     }
   }, [t, board_id, bulkDownload, dispatch]);
 
-  const skipEvent = useCallback((e: MouseEvent<HTMLDivElement>) => {
-    e.preventDefault();
-  }, []);
-
   const renderMenuFunc = useCallback(
     () => (
-      <InvMenuList visibility="visible" onContextMenu={skipEvent}>
+      <InvMenuList visibility="visible">
         <InvMenuGroup title={boardName}>
           <InvMenuItem
             icon={<FaPlus />}
@@ -131,7 +126,6 @@ const BoardContextMenu = ({
       isBulkDownloadEnabled,
       isSelectedForAutoAdd,
       setBoardToDelete,
-      skipEvent,
       t,
     ]
   );

invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/ImageContextMenu.tsx

@@ -2,7 +2,6 @@ import { useAppSelector } from 'app/store/storeHooks';
 import type { InvContextMenuProps } from 'common/components/InvContextMenu/InvContextMenu';
 import { InvContextMenu } from 'common/components/InvContextMenu/InvContextMenu';
 import { InvMenuList } from 'common/components/InvMenu/InvMenuList';
-import type { MouseEvent } from 'react';
 import { memo, useCallback } from 'react';
 import type { ImageDTO } from 'services/api/types';
 
@@ -17,29 +16,25 @@ type Props = {
 const ImageContextMenu = ({ imageDTO, children }: Props) => {
   const selectionCount = useAppSelector((s) => s.gallery.selection.length);
 
-  const skipEvent = useCallback((e: MouseEvent<HTMLDivElement>) => {
-    e.preventDefault();
-  }, []);
-
   const renderMenuFunc = useCallback(() => {
     if (!imageDTO) {
       return null;
     }
 
     if (selectionCount > 1) {
       return (
-        <InvMenuList visibility="visible" onContextMenu={skipEvent}>
+        <InvMenuList visibility="visible">
           <MultipleSelectionMenuItems />
         </InvMenuList>
       );
     }
 
     return (
-      <InvMenuList visibility="visible" onContextMenu={skipEvent}>
+      <InvMenuList visibility="visible">
         <SingleSelectionMenuItems imageDTO={imageDTO} />
       </InvMenuList>
     );
-  }, [imageDTO, selectionCount, skipEvent]);
+  }, [imageDTO, selectionCount]);
 
   return (
     <InvContextMenu renderMenu={renderMenuFunc}>{children}</InvContextMenu>

invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx

@@ -1,6 +1,6 @@
 import { useToken } from '@chakra-ui/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { useGlobalMenuCloseTrigger } from 'common/hooks/useGlobalMenuCloseTrigger';
+import { useGlobalMenuClose } from 'common/hooks/useGlobalMenuClose';
 import { useIsValidConnection } from 'features/nodes/hooks/useIsValidConnection';
 import { $mouseOverNode } from 'features/nodes/hooks/useMouseOverNode';
 import { useWorkflowWatcher } from 'features/nodes/hooks/useWorkflowWatcher';
@@ -158,7 +158,7 @@ export const Flow = memo(() => {
     [dispatch]
   );
 
-  const { onCloseGlobal } = useGlobalMenuCloseTrigger();
+  const { onCloseGlobal } = useGlobalMenuClose();
   const handlePaneClick = useCallback(() => {
     onCloseGlobal();
   }, [onCloseGlobal]);

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/FieldContextMenu.tsx

@@ -13,7 +13,7 @@ import {
   workflowExposedFieldAdded,
   workflowExposedFieldRemoved,
 } from 'features/nodes/store/workflowSlice';
-import type { MouseEvent, ReactNode } from 'react';
+import type { ReactNode } from 'react';
 import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { FaMinus, FaPlus } from 'react-icons/fa';
@@ -32,10 +32,6 @@ const FieldContextMenu = ({ nodeId, fieldName, kind, children }: Props) => {
   const input = useFieldInputKind(nodeId, fieldName);
   const { t } = useTranslation();
 
-  const skipEvent = useCallback((e: MouseEvent<HTMLDivElement>) => {
-    e.preventDefault();
-  }, []);
-
   const selectIsExposed = useMemo(
     () =>
       createSelector(selectWorkflowSlice, (workflow) => {
@@ -101,15 +97,15 @@ const FieldContextMenu = ({ nodeId, fieldName, kind, children }: Props) => {
   const renderMenuFunc = useCallback(
     () =>
       !menuItems.length ? null : (
-        <InvMenuList visibility="visible" onContextMenu={skipEvent}>
+        <InvMenuList visibility="visible">
           <InvMenuGroup
             title={label || fieldTemplateTitle || t('nodes.unknownField')}
           >
             {menuItems}
           </InvMenuGroup>
         </InvMenuList>
       ),
-    [fieldTemplateTitle, label, menuItems, skipEvent, t]
+    [fieldTemplateTitle, label, menuItems, t]
   );
 
   return (

invokeai/frontend/web/src/features/nodes/components/flow/nodes/common/NodeWrapper.tsx

@@ -3,7 +3,7 @@ import { Box, useToken } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import NodeSelectionOverlay from 'common/components/NodeSelectionOverlay';
-import { useGlobalMenuCloseTrigger } from 'common/hooks/useGlobalMenuCloseTrigger';
+import { useGlobalMenuClose } from 'common/hooks/useGlobalMenuClose';
 import { useMouseOverNode } from 'features/nodes/hooks/useMouseOverNode';
 import {
   nodeExclusivelySelected,
@@ -50,7 +50,7 @@ const NodeWrapper = (props: NodeWrapperProps) => {
   const dispatch = useAppDispatch();
 
   const opacity = useAppSelector((s) => s.nodes.nodeOpacity);
-  const { onCloseGlobal } = useGlobalMenuCloseTrigger();
+  const { onCloseGlobal } = useGlobalMenuClose();
 
   const handleClick = useCallback(
     (e: MouseEvent<HTMLDivElement>) => {

invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsMenu.tsx

@@ -7,7 +7,7 @@ import {
   InvMenuButton,
   InvMenuGroup,
 } from 'common/components/InvMenu/wrapper';
-import { useGlobalMenuCloseTrigger } from 'common/hooks/useGlobalMenuCloseTrigger';
+import { useGlobalMenuClose } from 'common/hooks/useGlobalMenuClose';
 import HotkeysModal from 'features/system/components/HotkeysModal/HotkeysModal';
 import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
 import { memo } from 'react';
@@ -23,7 +23,7 @@ import SettingsModal from './SettingsModal';
 const SettingsMenu = () => {
   const { t } = useTranslation();
   const { isOpen, onOpen, onClose } = useDisclosure();
-  useGlobalMenuCloseTrigger(onClose);
+  useGlobalMenuClose(onClose);
 
   const isBugLinkEnabled = useFeatureStatus('bugLink').isFeatureEnabled;
   const isDiscordLinkEnabled = useFeatureStatus('discordLink').isFeatureEnabled;