fix(ui): fix context menu on workflow editor

There is a tricky mouse event interaction between chakra's `useOutsideClick()` hook (used by chakra `<Menu />`) and reactflow. The hook doesn't work when you click the main reactflow area.

To get around this, I've used a dirty hack, copy-pasting the simple context menu component we use, and extending it slightly to respond to a global `contextMenusClosed` redux action.

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

@@ -0,0 +1,126 @@
+/**
+ * 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`
+ * - `useEffect()` to close the menu when `globalContextMenuCloseTrigger` changes
+ */
+
+import {
+  Menu,
+  MenuButton,
+  MenuButtonProps,
+  MenuProps,
+  Portal,
+  PortalProps,
+  useEventListener,
+} from '@chakra-ui/react';
+import { useAppSelector } from 'app/store/storeHooks';
+import * as React from 'react';
+import {
+  MutableRefObject,
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+} from 'react';
+
+export interface IAIContextMenuProps<T extends HTMLElement> {
+  renderMenu: () => JSX.Element | null;
+  children: (ref: MutableRefObject<T | null>) => JSX.Element | null;
+  menuProps?: Omit<MenuProps, 'children'> & { children?: React.ReactNode };
+  portalProps?: Omit<PortalProps, 'children'> & { children?: React.ReactNode };
+  menuButtonProps?: MenuButtonProps;
+}
+
+export function IAIContextMenu<T extends HTMLElement = HTMLElement>(
+  props: IAIContextMenuProps<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 targetRef = useRef<T>(null);
+
+  const globalContextMenuCloseTrigger = useAppSelector(
+    (state) => state.ui.globalContextMenuCloseTrigger
+  );
+
+  useEffect(() => {
+    if (isOpen) {
+      setTimeout(() => {
+        setIsRendered(true);
+        setTimeout(() => {
+          setIsDeferredOpen(true);
+        });
+      });
+    } else {
+      setIsDeferredOpen(false);
+      const timeout = setTimeout(() => {
+        setIsRendered(isOpen);
+      }, 1000);
+      return () => clearTimeout(timeout);
+    }
+  }, [isOpen]);
+
+  useEffect(() => {
+    setIsOpen(false);
+    setIsDeferredOpen(false);
+    setIsRendered(false);
+  }, [globalContextMenuCloseTrigger]);
+
+  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);
+    }
+  });
+
+  const onCloseHandler = useCallback(() => {
+    props.menuProps?.onClose?.();
+    setIsOpen(false);
+  }, [props.menuProps]);
+
+  return (
+    <>
+      {props.children(targetRef)}
+      {isRendered && (
+        <Portal {...props.portalProps}>
+          <Menu
+            isOpen={isDeferredOpen}
+            gutter={0}
+            {...props.menuProps}
+            onClose={onCloseHandler}
+          >
+            <MenuButton
+              aria-hidden={true}
+              w={1}
+              h={1}
+              style={{
+                position: 'absolute',
+                left: position[0],
+                top: position[1],
+                cursor: 'default',
+              }}
+              {...props.menuButtonProps}
+            />
+            {props.renderMenu()}
+          </Menu>
+        </Portal>
+      )}
+    </>
+  );
+}

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

@@ -16,6 +16,7 @@ import ImageContextMenu from 'features/gallery/components/ImageContextMenu/Image
 import {
   MouseEvent,
   ReactElement,
+  ReactNode,
   SyntheticEvent,
   memo,
   useCallback,
@@ -32,6 +33,17 @@ import {
   TypesafeDroppableData,
 } from 'features/dnd/types';
 
+const defaultUploadElement = (
+  <Icon
+    as={FaUpload}
+    sx={{
+      boxSize: 16,
+    }}
+  />
+);
+
+const defaultNoContentFallback = <IAINoContentFallback icon={FaImage} />;
+
 type IAIDndImageProps = FlexProps & {
   imageDTO: ImageDTO | undefined;
   onError?: (event: SyntheticEvent<HTMLImageElement>) => void;
@@ -47,13 +59,14 @@ type IAIDndImageProps = FlexProps & {
   fitContainer?: boolean;
   droppableData?: TypesafeDroppableData;
   draggableData?: TypesafeDraggableData;
-  dropLabel?: string;
+  dropLabel?: ReactNode;
   isSelected?: boolean;
   thumbnail?: boolean;
   noContentFallback?: ReactElement;
   useThumbailFallback?: boolean;
   withHoverOverlay?: boolean;
   children?: JSX.Element;
+  uploadElement?: ReactNode;
 };
 
 const IAIDndImage = (props: IAIDndImageProps) => {
@@ -74,7 +87,8 @@ const IAIDndImage = (props: IAIDndImageProps) => {
     dropLabel,
     isSelected = false,
     thumbnail = false,
-    noContentFallback = <IAINoContentFallback icon={FaImage} />,
+    noContentFallback = defaultNoContentFallback,
+    uploadElement = defaultUploadElement,
     useThumbailFallback,
     withHoverOverlay = false,
     children,
@@ -193,12 +207,7 @@ const IAIDndImage = (props: IAIDndImageProps) => {
                 {...getUploadButtonProps()}
               >
                 <input {...getUploadInputProps()} />
-                <Icon
+                {uploadElement}
-                  as={FaUpload}
-                  sx={{
-                    boxSize: 16,
-                  }}
-                />
               </Flex>
             </>
           )}
@@ -210,6 +219,7 @@ const IAIDndImage = (props: IAIDndImageProps) => {
               onClick={onClick}
             />
           )}
+          {children}
           {!isDropDisabled && (
             <IAIDroppable
               data={droppableData}
@@ -217,7 +227,6 @@ const IAIDndImage = (props: IAIDndImageProps) => {
               dropLabel={dropLabel}
             />
           )}
-          {children}
         </Flex>
       )}
     </ImageContextMenu>

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

@@ -1,5 +1,8 @@
 import { MenuList } from '@chakra-ui/react';
-import { ContextMenu, ContextMenuProps } from 'chakra-ui-contextmenu';
+import {
+  IAIContextMenu,
+  IAIContextMenuProps,
+} from 'common/components/IAIContextMenu';
 import { MouseEvent, memo, useCallback } from 'react';
 import { ImageDTO } from 'services/api/types';
 import { menuListMotionProps } from 'theme/components/menu';
@@ -12,7 +15,7 @@ import MultipleSelectionMenuItems from './MultipleSelectionMenuItems';
 
 type Props = {
   imageDTO: ImageDTO | undefined;
-  children: ContextMenuProps<HTMLDivElement>['children'];
+  children: IAIContextMenuProps<HTMLDivElement>['children'];
 };
 
 const selector = createSelector(
@@ -33,7 +36,7 @@ const ImageContextMenu = ({ imageDTO, children }: Props) => {
   }, []);
 
   return (
-    <ContextMenu<HTMLDivElement>
+    <IAIContextMenu<HTMLDivElement>
       menuProps={{ size: 'sm', isLazy: true }}
       menuButtonProps={{
         bg: 'transparent',
@@ -68,7 +71,7 @@ const ImageContextMenu = ({ imageDTO, children }: Props) => {
       }}
     >
       {children}
-    </ContextMenu>
+    </IAIContextMenu>
   );
 };
 

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

@@ -1,5 +1,6 @@
 import { useToken } from '@chakra-ui/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { contextMenusClosed } from 'features/ui/store/uiSlice';
 import { useCallback } from 'react';
 import {
   Background,
@@ -114,6 +115,10 @@ export const Flow = () => {
     [dispatch]
   );
 
+  const handlePaneClick = useCallback(() => {
+    dispatch(contextMenusClosed());
+  }, [dispatch]);
+
   return (
     <ReactFlow
       defaultViewport={viewport}
@@ -138,6 +143,7 @@ export const Flow = () => {
       connectionRadius={30}
       proOptions={proOptions}
       style={{ borderRadius }}
+      onPaneClick={handlePaneClick}
     >
       <TopLeftPanel />
       <TopCenterPanel />

invokeai/frontend/web/src/features/ui/store/uiPersistDenylist.ts

@@ -3,4 +3,7 @@ import { UIState } from './uiTypes';
 /**
  * UI slice persist denylist
  */
-export const uiPersistDenylist: (keyof UIState)[] = ['shouldShowImageDetails'];
+export const uiPersistDenylist: (keyof UIState)[] = [
+  'shouldShowImageDetails',
+  'globalContextMenuCloseTrigger',
+];

invokeai/frontend/web/src/features/ui/store/uiSlice.ts

@@ -20,6 +20,7 @@ export const initialUIState: UIState = {
   shouldShowProgressInViewer: true,
   shouldShowEmbeddingPicker: false,
   favoriteSchedulers: [],
+  globalContextMenuCloseTrigger: 0,
 };
 
 export const uiSlice = createSlice({
@@ -96,6 +97,9 @@ export const uiSlice = createSlice({
     toggleEmbeddingPicker: (state) => {
       state.shouldShowEmbeddingPicker = !state.shouldShowEmbeddingPicker;
     },
+    contextMenusClosed: (state) => {
+      state.globalContextMenuCloseTrigger += 1;
+    },
   },
   extraReducers(builder) {
     builder.addCase(initialImageChanged, (state) => {
@@ -122,6 +126,7 @@ export const {
   setShouldShowProgressInViewer,
   favoriteSchedulersChanged,
   toggleEmbeddingPicker,
+  contextMenusClosed,
 } = uiSlice.actions;
 
 export default uiSlice.reducer;

invokeai/frontend/web/src/features/ui/store/uiTypes.ts

@@ -26,4 +26,5 @@ export interface UIState {
   shouldShowProgressInViewer: boolean;
   shouldShowEmbeddingPicker: boolean;
   favoriteSchedulers: SchedulerParam[];
+  globalContextMenuCloseTrigger: number;
 }