From 61060f032a1fa88bf0f971b348225e74625744ce Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 6 Dec 2023 18:54:23 +1100
Subject: [PATCH] feat(ui): abstract out the global menu close trigger

This logic is moved into a hook.

This is needed for our context menus to close when the user clicks something in reactflow. It needed to be extended to support menus also.
---
 .../src/common/components/IAIContextMenu.tsx  | 11 +++-----
 .../common/hooks/useGlobalMenuCloseTrigger.ts | 25 +++++++++++++++++++
 .../features/nodes/components/flow/Flow.tsx   |  4 +--
 .../flow/nodes/common/NodeWrapper.tsx         |  4 +--
 .../features/system/components/SiteHeader.tsx |  6 ++++-
 .../web/src/features/ui/store/uiSlice.ts      |  8 +++---
 .../web/src/features/ui/store/uiTypes.ts      |  2 +-
 7 files changed, 43 insertions(+), 17 deletions(-)
 create mode 100644 invokeai/frontend/web/src/common/hooks/useGlobalMenuCloseTrigger.ts

diff --git a/invokeai/frontend/web/src/common/components/IAIContextMenu.tsx b/invokeai/frontend/web/src/common/components/IAIContextMenu.tsx
index 757faca866..9fb6f1b835 100644
--- a/invokeai/frontend/web/src/common/components/IAIContextMenu.tsx
+++ b/invokeai/frontend/web/src/common/components/IAIContextMenu.tsx
@@ -22,7 +22,7 @@ import {
   PortalProps,
   useEventListener,
 } from '@chakra-ui/react';
-import { useAppSelector } from 'app/store/storeHooks';
+import { useGlobalMenuCloseTrigger } from 'common/hooks/useGlobalMenuCloseTrigger';
 import * as React from 'react';
 import {
   MutableRefObject,
@@ -49,10 +49,6 @@ export function IAIContextMenu<T extends HTMLElement = HTMLElement>(
   const [position, setPosition] = useState<[number, number]>([0, 0]);
   const targetRef = useRef<T>(null);
 
-  const globalContextMenuCloseTrigger = useAppSelector(
-    (state) => state.ui.globalContextMenuCloseTrigger
-  );
-
   useEffect(() => {
     if (isOpen) {
       setTimeout(() => {
@@ -70,11 +66,12 @@ export function IAIContextMenu<T extends HTMLElement = HTMLElement>(
     }
   }, [isOpen]);
 
-  useEffect(() => {
+  const onClose = useCallback(() => {
     setIsOpen(false);
     setIsDeferredOpen(false);
     setIsRendered(false);
-  }, [globalContextMenuCloseTrigger]);
+  }, []);
+  useGlobalMenuCloseTrigger(onClose);
 
   useEventListener('contextmenu', (e) => {
     if (
diff --git a/invokeai/frontend/web/src/common/hooks/useGlobalMenuCloseTrigger.ts b/invokeai/frontend/web/src/common/hooks/useGlobalMenuCloseTrigger.ts
new file mode 100644
index 0000000000..0fd404fdc3
--- /dev/null
+++ b/invokeai/frontend/web/src/common/hooks/useGlobalMenuCloseTrigger.ts
@@ -0,0 +1,25 @@
+import { useAppSelector } from 'app/store/storeHooks';
+import { useEffect } from 'react';
+
+/**
+ * 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 `globalMenuCloseTrigger: number` in `ui` slice
+ * - increment it in `onPaneClick`
+ * - `useEffect()` to close the menu when `globalMenuCloseTrigger` changes
+ */
+
+export const useGlobalMenuCloseTrigger = (onClose: () => void) => {
+  const globalMenuCloseTrigger = useAppSelector(
+    (state) => state.ui.globalMenuCloseTrigger
+  );
+
+  useEffect(() => {
+    onClose();
+  }, [globalMenuCloseTrigger, onClose]);
+};
diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
index ea83a540a9..2f0695f03a 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
@@ -4,7 +4,7 @@ import { stateSelector } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import { $flow } from 'features/nodes/store/reactFlowInstance';
-import { contextMenusClosed } from 'features/ui/store/uiSlice';
+import { bumpGlobalMenuCloseTrigger } from 'features/ui/store/uiSlice';
 import { MouseEvent, useCallback, useRef } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import {
@@ -153,7 +153,7 @@ export const Flow = () => {
   );
 
   const handlePaneClick = useCallback(() => {
-    dispatch(contextMenusClosed());
+    dispatch(bumpGlobalMenuCloseTrigger());
   }, [dispatch]);
 
   const onInit: OnInit = useCallback((flow) => {
diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/nodes/common/NodeWrapper.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/nodes/common/NodeWrapper.tsx
index b6ccd4ae9f..155e95da94 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/nodes/common/NodeWrapper.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/nodes/common/NodeWrapper.tsx
@@ -15,7 +15,7 @@ import {
   NODE_WIDTH,
 } from 'features/nodes/types/constants';
 import { zNodeStatus } from 'features/nodes/types/invocation';
-import { contextMenusClosed } from 'features/ui/store/uiSlice';
+import { bumpGlobalMenuCloseTrigger } from 'features/ui/store/uiSlice';
 import {
   MouseEvent,
   PropsWithChildren,
@@ -70,7 +70,7 @@ const NodeWrapper = (props: NodeWrapperProps) => {
       if (!e.ctrlKey && !e.altKey && !e.metaKey && !e.shiftKey) {
         dispatch(nodeExclusivelySelected(nodeId));
       }
-      dispatch(contextMenusClosed());
+      dispatch(bumpGlobalMenuCloseTrigger());
     },
     [dispatch, nodeId]
   );
diff --git a/invokeai/frontend/web/src/features/system/components/SiteHeader.tsx b/invokeai/frontend/web/src/features/system/components/SiteHeader.tsx
index 5057af8dab..6bc8ae6f59 100644
--- a/invokeai/frontend/web/src/features/system/components/SiteHeader.tsx
+++ b/invokeai/frontend/web/src/features/system/components/SiteHeader.tsx
@@ -6,6 +6,7 @@ import {
   MenuItem,
   MenuList,
   Spacer,
+  useDisclosure,
 } from '@chakra-ui/react';
 import IAIIconButton from 'common/components/IAIIconButton';
 import { memo } from 'react';
@@ -24,9 +25,12 @@ import HotkeysModal from './HotkeysModal/HotkeysModal';
 import InvokeAILogoComponent from './InvokeAILogoComponent';
 import SettingsModal from './SettingsModal/SettingsModal';
 import StatusIndicator from './StatusIndicator';
+import { useGlobalMenuCloseTrigger } from 'common/hooks/useGlobalMenuCloseTrigger';
 
 const SiteHeader = () => {
   const { t } = useTranslation();
+  const { isOpen, onOpen, onClose } = useDisclosure();
+  useGlobalMenuCloseTrigger(onClose);
 
   const isBugLinkEnabled = useFeatureStatus('bugLink').isFeatureEnabled;
   const isDiscordLinkEnabled = useFeatureStatus('discordLink').isFeatureEnabled;
@@ -46,7 +50,7 @@ const SiteHeader = () => {
       <Spacer />
       <StatusIndicator />
 
-      <Menu>
+      <Menu isOpen={isOpen} onOpen={onOpen} onClose={onClose}>
         <MenuButton
           as={IAIIconButton}
           variant="link"
diff --git a/invokeai/frontend/web/src/features/ui/store/uiSlice.ts b/invokeai/frontend/web/src/features/ui/store/uiSlice.ts
index 69cfe42827..a5ecfc34c2 100644
--- a/invokeai/frontend/web/src/features/ui/store/uiSlice.ts
+++ b/invokeai/frontend/web/src/features/ui/store/uiSlice.ts
@@ -16,7 +16,7 @@ export const initialUIState: UIState = {
   shouldShowEmbeddingPicker: false,
   shouldAutoChangeDimensions: false,
   favoriteSchedulers: [],
-  globalContextMenuCloseTrigger: 0,
+  globalMenuCloseTrigger: 0,
   panels: {},
 };
 
@@ -60,8 +60,8 @@ export const uiSlice = createSlice({
     setShouldAutoChangeDimensions: (state, action: PayloadAction<boolean>) => {
       state.shouldAutoChangeDimensions = action.payload;
     },
-    contextMenusClosed: (state) => {
-      state.globalContextMenuCloseTrigger += 1;
+    bumpGlobalMenuCloseTrigger: (state) => {
+      state.globalMenuCloseTrigger += 1;
     },
     panelsChanged: (
       state,
@@ -88,7 +88,7 @@ export const {
   favoriteSchedulersChanged,
   toggleEmbeddingPicker,
   setShouldAutoChangeDimensions,
-  contextMenusClosed,
+  bumpGlobalMenuCloseTrigger,
   panelsChanged,
 } = uiSlice.actions;
 
diff --git a/invokeai/frontend/web/src/features/ui/store/uiTypes.ts b/invokeai/frontend/web/src/features/ui/store/uiTypes.ts
index b532043054..1a25d4d3d6 100644
--- a/invokeai/frontend/web/src/features/ui/store/uiTypes.ts
+++ b/invokeai/frontend/web/src/features/ui/store/uiTypes.ts
@@ -24,6 +24,6 @@ export interface UIState {
   shouldShowEmbeddingPicker: boolean;
   shouldAutoChangeDimensions: boolean;
   favoriteSchedulers: ParameterScheduler[];
-  globalContextMenuCloseTrigger: number;
+  globalMenuCloseTrigger: number;
   panels: Record<string, string>;
 }