From 6209fef63d7b51f74f9fd14a5678ceadc279e167 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Sun, 31 Dec 2023 19:34:05 +1100
Subject: [PATCH] fix(ui): focus add node popover on open

Need an extra ref to pass to the InvSelect component.
---
 .../web/src/common/components/InvSelect/InvSelect.tsx | 11 +++++++++--
 .../web/src/common/components/InvSelect/types.ts      |  1 +
 .../components/flow/AddNodePopover/AddNodePopover.tsx |  4 ++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/invokeai/frontend/web/src/common/components/InvSelect/InvSelect.tsx b/invokeai/frontend/web/src/common/components/InvSelect/InvSelect.tsx
index 69705f588d..6d44e6bfed 100644
--- a/invokeai/frontend/web/src/common/components/InvSelect/InvSelect.tsx
+++ b/invokeai/frontend/web/src/common/components/InvSelect/InvSelect.tsx
@@ -4,7 +4,7 @@ import type {
   StylesConfig,
 } from 'chakra-react-select';
 import { Select as ChakraReactSelect } from 'chakra-react-select';
-import { memo, useMemo } from 'react';
+import { memo, useEffect, useMemo } from 'react';
 
 import { CustomMenuList } from './CustomMenuList';
 import { CustomOption } from './CustomOption';
@@ -29,7 +29,7 @@ const components: SelectComponentsConfig<
 };
 
 export const InvSelect = memo((props: InvSelectProps) => {
-  const { sx, selectRef, ...rest } = props;
+  const { sx, selectRef, inputRef, ...rest } = props;
   const chakraStyles = useMemo<CustomChakraStylesConfig>(
     () => ({
       container: (provided, _state) => ({ ...provided, w: 'full', ...sx }),
@@ -55,6 +55,13 @@ export const InvSelect = memo((props: InvSelectProps) => {
     [sx]
   );
 
+  useEffect(() => {
+    if (!inputRef) {
+      return;
+    }
+    inputRef.current = selectRef?.current?.inputRef ?? null;
+  }, [inputRef, selectRef]);
+
   return (
     <ChakraReactSelect<InvSelectOption, false, GroupBase<InvSelectOption>>
       ref={selectRef}
diff --git a/invokeai/frontend/web/src/common/components/InvSelect/types.ts b/invokeai/frontend/web/src/common/components/InvSelect/types.ts
index 9f05ca0763..c52615c3fb 100644
--- a/invokeai/frontend/web/src/common/components/InvSelect/types.ts
+++ b/invokeai/frontend/web/src/common/components/InvSelect/types.ts
@@ -28,6 +28,7 @@ export type InvSelectProps = ChakraReactSelectProps<
   selectRef?: React.RefObject<
     SelectInstance<InvSelectOption, false, GroupBase<InvSelectOption>>
   >;
+  inputRef?: React.MutableRefObject<HTMLInputElement | null>;
 };
 export type CustomChakraStylesConfig = ChakraStylesConfig<
   InvSelectOption,
diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/AddNodePopover/AddNodePopover.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/AddNodePopover/AddNodePopover.tsx
index f7e4458799..5139dc1ff8 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/AddNodePopover/AddNodePopover.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/AddNodePopover/AddNodePopover.tsx
@@ -65,6 +65,8 @@ const AddNodePopover = () => {
   const toaster = useAppToaster();
   const { t } = useTranslation();
   const selectRef = useRef<SelectInstance<InvSelectOption> | null>(null);
+  const inputRef = useRef<HTMLInputElement>(null);
+
   const fieldFilter = useAppSelector(
     (state) => state.nodes.connectionStartFieldType
   );
@@ -201,6 +203,7 @@ const AddNodePopover = () => {
       closeDelay={0}
       closeOnBlur={true}
       returnFocusOnClose={true}
+      initialFocusRef={inputRef}
     >
       <InvPopoverAnchor>
         <Flex
@@ -230,6 +233,7 @@ const AddNodePopover = () => {
             onChange={onChange}
             onMenuClose={onClose}
             onKeyDown={onKeyDown}
+            inputRef={inputRef}
           />
         </InvPopoverBody>
       </InvPopoverContent>