fix: focus being lost when input value is set from external state (#51)

Author: BramKaashoek

.changeset/custom-components-story.md

@@ -0,0 +1,5 @@
+---
+"react-loqate": patch
+---
+
+- Fix focus being lost when input value is set from external state

src/index.test.tsx

@@ -378,6 +378,64 @@ it('accepts origin and bias options', async () => {
   });
 });
 
+it('preserves focus when using custom Input with external state management', async () => {
+  function TestComponent() {
+    const [, setExternalState] = React.useState('');
+
+    return (
+      <AddressSearch
+        locale="en-GB"
+        apiKey="test-key"
+        onSelect={() => {}}
+        components={{
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          Input: React.forwardRef<HTMLInputElement, any>(
+            ({ value, onChange, ...rest }, ref) => {
+              React.useEffect(() => {
+                setExternalState(value || '');
+              }, [value]);
+
+              return (
+                <input
+                  ref={ref}
+                  value={value || ''}
+                  onChange={(event) => {
+                    onChange?.(event);
+                    setExternalState(event.target.value);
+                  }}
+                  {...rest}
+                  data-testid="external-state-input"
+                />
+              );
+            }
+          ),
+        }}
+      />
+    );
+  }
+
+  render(<TestComponent />);
+
+  const input = screen.getByTestId('external-state-input') as HTMLInputElement;
+  input.focus();
+  expect(document.activeElement).toBe(input);
+
+  fireEvent.change(input, { target: { value: 'a' } });
+
+  await screen.findByRole('list');
+
+  await waitFor(
+    () => {
+      const currentInput = screen.getByTestId('external-state-input');
+      expect(document.activeElement).toBe(currentInput);
+    },
+    { timeout: 1000 }
+  );
+
+  const suggestions = screen.getAllByRole('listitem');
+  expect(suggestions.length).toBeGreaterThan(0);
+});
+
 it('disables browser autocomplete by default', () => {
   render(<AddressSearch locale="en-GB" apiKey="some-key" onSelect={vi.fn()} />);
 

src/index.tsx

@@ -3,7 +3,6 @@ import React, {
   type ChangeEvent,
   type ComponentType,
   useMemo,
-  useRef,
   useState,
 } from 'react';
 import DefaultInput from './components/DefaultInput';
@@ -13,6 +12,7 @@ import ClickAwayListener from './utils/ClickAwayListener';
 import Loqate from './utils/Loqate';
 import Portal from './utils/Portal';
 import useDebounceEffect from './utils/useDebounceEffect';
+import usePreserveFocus from './utils/usePreserveFocus';
 
 export interface Props {
   locale: string;
@@ -143,7 +143,9 @@ function AddressSearch(props: Props): JSX.Element {
   const [value, setValue] = useState('');
   const [, setError] = useState(null);
 
-  const anchorRef = useRef<HTMLInputElement>(null);
+  const { elementRef: anchorRef, preserveFocus } =
+    usePreserveFocus<HTMLInputElement>();
+
   const rect = anchorRef.current?.getBoundingClientRect();
 
   async function find(text: string, containerId?: string): Promise<Item[]> {
@@ -202,6 +204,8 @@ function AddressSearch(props: Props): JSX.Element {
   }: ChangeEvent<HTMLInputElement>): Promise<void> {
     const { value: search } = target;
 
+    // Custom Input components with external state management can cause DOM reconciliation issues that lose focus
+    preserveFocus();
     setValue(search);
   }
 

src/stories/AddressSearch.stories.tsx

@@ -1,6 +1,6 @@
-import { Meta, StoryObj } from '@storybook/react';
+import type { Meta, StoryObj } from '@storybook/react';
 import React, { useState } from 'react';
-import AddressSearch, { Address, Props } from '..';
+import AddressSearch, { type Address, type Props } from '..';
 
 const meta: Meta = {
   title: 'Loqate Address Search',

src/stories/CustomInput.stories.tsx

@@ -0,0 +1,128 @@
+import type { Meta, StoryObj } from '@storybook/react';
+import React, { forwardRef, useState } from 'react';
+import AddressSearch, { type Address, type Item } from '..';
+
+const meta: Meta = {
+  title: 'Custom Components',
+  component: AddressSearch,
+};
+
+export default meta;
+
+type Story = StoryObj<typeof AddressSearch>;
+
+const CustomInput = forwardRef<HTMLInputElement, React.ComponentProps<'input'>>(
+  ({ value, onChange, ...rest }, ref) => {
+    return (
+      <input
+        ref={ref}
+        value={value || ''}
+        onChange={onChange}
+        style={{
+          padding: '12px 16px',
+          border: '2px solid #3b82f6',
+          borderRadius: '8px',
+          fontSize: '16px',
+          width: '100%',
+          outline: 'none',
+        }}
+        placeholder="Enter your address"
+        {...rest}
+      />
+    );
+  }
+);
+
+const CustomList = forwardRef<HTMLUListElement, React.ComponentProps<'ul'>>(
+  (props, ref) => {
+    return (
+      <ul
+        ref={ref}
+        {...props}
+        style={{
+          ...props.style,
+          backgroundColor: '#f8fafc',
+          border: '2px solid #3b82f6',
+          borderRadius: '8px',
+          listStyle: 'none',
+          margin: 0,
+          padding: 0,
+        }}
+      />
+    );
+  }
+);
+
+const CustomListItem = forwardRef<
+  HTMLLIElement,
+  React.ComponentProps<'li'> & { suggestion: Item; value?: string }
+>((props, ref) => {
+  const {
+    suggestion,
+    onClick,
+    onKeyDown,
+    onMouseEnter,
+    onMouseLeave,
+    style,
+    ...rest
+  } = props;
+  return (
+    <li
+      ref={ref}
+      {...rest}
+      onClick={onClick}
+      onKeyDown={(e) => {
+        if (e.key === 'Enter' || e.key === ' ') {
+          e.preventDefault();
+          onClick?.(e as unknown as React.MouseEvent<HTMLLIElement>);
+        }
+        onKeyDown?.(e);
+      }}
+      style={{
+        padding: '12px 16px',
+        borderBottom: '1px solid #e2e8f0',
+        cursor: 'pointer',
+        backgroundColor: 'white',
+        ...style,
+      }}
+      onMouseEnter={(e) => {
+        e.currentTarget.style.backgroundColor = '#eff6ff';
+        onMouseEnter?.(e);
+      }}
+      onMouseLeave={(e) => {
+        e.currentTarget.style.backgroundColor = 'white';
+        onMouseLeave?.(e);
+      }}
+    >
+      {suggestion.Text} {suggestion.Description}
+    </li>
+  );
+});
+
+const Template = (props: React.ComponentProps<typeof AddressSearch>) => {
+  const [result, setResult] = useState<Address>();
+
+  return (
+    <>
+      <AddressSearch {...props} onSelect={setResult} />
+      <pre>{JSON.stringify(result, null, 2)}</pre>
+    </>
+  );
+};
+
+export const CustomComponents: Story = {
+  args: {
+    // @ts-expect-error: env does exist
+    // We need to prefix with STORYBOOK, otherwise Storybook will ignore the variable
+    apiKey: import.meta.env.STORYBOOK_API_KEY ?? '',
+    countries: ['US', 'GB'],
+    locale: 'en_US',
+    inline: true,
+    components: {
+      Input: CustomInput,
+      List: CustomList,
+      ListItem: CustomListItem,
+    },
+  },
+  render: Template,
+};

src/utils/usePreserveFocus.test.tsx

@@ -0,0 +1,56 @@
+import { act, render, waitFor } from '@testing-library/react';
+import React, { useState } from 'react';
+import { describe, expect, it, vi } from 'vitest';
+import usePreserveFocus from './usePreserveFocus';
+
+describe('usePreserveFocus', () => {
+  it('calls focus() asynchronously via Promise.resolve when element is focused', async () => {
+    const mockFocus = vi.fn();
+
+    function TestComponent() {
+      const [count, setCount] = useState(0);
+      const { elementRef, preserveFocus } =
+        usePreserveFocus<HTMLInputElement>();
+
+      React.useLayoutEffect(() => {
+        if (elementRef.current) {
+          elementRef.current.focus = mockFocus;
+          Object.defineProperty(document, 'activeElement', {
+            value: elementRef.current,
+            writable: true,
+            configurable: true,
+          });
+        }
+      });
+
+      const handleClick = () => {
+        preserveFocus();
+        setCount(count + 1);
+      };
+
+      return (
+        <div>
+          <input ref={elementRef} data-testid="test-input" />
+          <button
+            type="button"
+            onClick={handleClick}
+            data-testid="trigger-action"
+          >
+            Count: {count}
+          </button>
+        </div>
+      );
+    }
+
+    const { getByTestId } = render(<TestComponent />);
+    const button = getByTestId('trigger-action');
+
+    act(() => {
+      button.click();
+    });
+
+    await waitFor(() => {
+      expect(mockFocus).toHaveBeenCalledTimes(1);
+    });
+  });
+});

src/utils/usePreserveFocus.ts

@@ -0,0 +1,34 @@
+import { useLayoutEffect, useRef, useState } from 'react';
+
+function usePreserveFocus<T extends HTMLElement = HTMLElement>(): {
+  elementRef: React.RefObject<T>;
+  preserveFocus: () => void;
+} {
+  const elementRef = useRef<T>(null);
+  const [shouldRestoreFocus, setShouldRestoreFocus] = useState(false);
+
+  const preserveFocus = () => {
+    if (elementRef.current && document.activeElement === elementRef.current) {
+      setShouldRestoreFocus(true);
+    }
+  };
+
+  useLayoutEffect(() => {
+    if (shouldRestoreFocus && elementRef.current) {
+      // Use a microtask to ensure this runs after any DOM updates
+      Promise.resolve().then(() => {
+        if (elementRef.current) {
+          elementRef.current.focus();
+        }
+      });
+      setShouldRestoreFocus(false);
+    }
+  });
+
+  return {
+    elementRef,
+    preserveFocus,
+  };
+}
+
+export default usePreserveFocus;