fix(snaps): Keep focus on input if interface re-renders

ui/components/app/snaps/snap-ui-input/snap-ui-input.tsx

@@ -2,6 +2,7 @@ import React, {
   ChangeEvent,
   FunctionComponent,
   useEffect,
+  useRef,
   useState,
 } from 'react';
 import { useSnapInterfaceContext } from '../../../../contexts/snaps';
@@ -15,7 +16,10 @@ export type SnapUIInputProps = {
 export const SnapUIInput: FunctionComponent<
   SnapUIInputProps & FormTextFieldProps<'div'>
 > = ({ name, form, ...props }) => {
-  const { handleInputChange, getValue } = useSnapInterfaceContext();
+  const { handleInputChange, getValue, focusedInput, setCurrentFocusedInput } =
+    useSnapInterfaceContext();
+
+  const inputRef = useRef<HTMLDivElement>(null);
 
   const initialValue = getValue(name, form) as string;
 
@@ -27,14 +31,29 @@ export const SnapUIInput: FunctionComponent<
     }
   }, [initialValue]);
 
+  /*
+   * Focus input if the last focused input was this input
+   * This avoids loosing the focus when the UI is re-rendered
+   */
+  useEffect(() => {
+    if (inputRef.current && name === focusedInput) {
+      (inputRef.current.children[0] as HTMLInputElement).focus();
+    }
+  }, [inputRef]);
+
   const handleChange = (event: ChangeEvent<HTMLInputElement>) => {
     setValue(event.target.value);
     handleInputChange(name, event.target.value ?? null, form);
   };
 
+  const handleFocus = () => setCurrentFocusedInput(name);
+  const handleBlur = () => setCurrentFocusedInput(null);
+
   return (
     <FormTextField
-      autoFocus
+      ref={inputRef}
+      onFocus={handleFocus}
+      onBlur={handleBlur}
       className="snap-ui-renderer__input"
       id={name}
       value={value}

ui/contexts/snaps/snap-interface.tsx

@@ -42,11 +42,15 @@ export type HandleFileChange = (
   form?: string,
 ) => void;
 
+export type SetCurrentInputFocus = (name: string | null) => void;
+
 export type SnapInterfaceContextType = {
   handleEvent: HandleEvent;
   getValue: GetValue;
   handleInputChange: HandleInputChange;
   handleFileChange: HandleFileChange;
+  setCurrentFocusedInput: SetCurrentInputFocus;
+  focusedInput: string | null;
   snapId: string;
 };
 
@@ -80,6 +84,7 @@ export const SnapInterfaceContextProvider: FunctionComponent<
   // UI. It's kept in a ref to avoid useless re-rendering of the entire tree of
   // components.
   const internalState = useRef<InterfaceState>(initialState ?? {});
+  const focusedInput = useRef<string | null>(null);
 
   // Since the internal state is kept in a reference, it won't update when the
   // interface is updated. We have to manually update it.
@@ -237,13 +242,18 @@ export const SnapInterfaceContextProvider: FunctionComponent<
     return undefined;
   };
 
+  const setCurrentFocusedInput: SetCurrentInputFocus = (name) =>
+    (focusedInput.current = name);
+
   return (
     <SnapInterfaceContext.Provider
       value={{
         handleEvent,
         getValue,
         handleInputChange,
         handleFileChange,
+        setCurrentFocusedInput,
+        focusedInput: focusedInput.current,
         snapId,
       }}
     >