DebouncedInput: support aggressive focus retention

src/apps/chat/components/layout-bar/useLLMDropdown.tsx

@@ -8,6 +8,7 @@ import SettingsIcon from '@mui/icons-material/Settings';
 import { DLLM, DLLMId, DModelSourceId, useModelsStore } from '~/modules/llms/store-llms';
 import { findVendorById } from '~/modules/llms/vendors/vendors.registry';
 
+import { DebouncedInputMemo } from '~/common/components/DebouncedInput';
 import { DropdownItems, PageBarDropdownMemo } from '~/common/layout/optima/components/PageBarDropdown';
 import { GoodTooltip } from '~/common/components/GoodTooltip';
 import { KeyStroke } from '~/common/components/KeyStroke';

src/apps/chat/components/layout-drawer/ChatDrawer.tsx

@@ -17,7 +17,7 @@ import StarOutlineRoundedIcon from '@mui/icons-material/StarOutlineRounded';
 import type { DConversationId } from '~/common/stores/chat/chat.conversation';
 import { CloseableMenu } from '~/common/components/CloseableMenu';
 import { DFolder, useFolderStore } from '~/common/state/store-folders';
-import { DebounceInputMemo } from '~/common/components/DebounceInput';
+import { DebouncedInputMemo } from '~/common/components/DebouncedInput';
 import { FoldersToggleOff } from '~/common/components/icons/FoldersToggleOff';
 import { FoldersToggleOn } from '~/common/components/icons/FoldersToggleOn';
 import { PageDrawerHeader } from '~/common/layout/optima/components/PageDrawerHeader';
@@ -280,7 +280,7 @@ function ChatDrawer(props: {
       <Box sx={{ display: 'flex', flexDirection: 'column', m: 2, gap: 2 }}>
 
         {/* Search Input Field */}
-        <DebounceInputMemo
+        <DebouncedInputMemo
           minChars={2}
           onDebounce={setDebouncedSearchQuery}
           debounceTimeout={300}

src/common/components/DebounceInput.tsx

@@ -1,73 +0,0 @@
-import * as React from 'react';
-
-import type { InputProps } from '@mui/joy/Input';
-import { Box, IconButton, Input } from '@mui/joy';
-import ClearIcon from '@mui/icons-material/Clear';
-import SearchIcon from '@mui/icons-material/Search';
-
-type DebounceInputProps = Omit<InputProps, 'onChange'> & {
-  minChars?: number;
-  onDebounce: (value: string) => void;
-  debounceTimeout: number;
-};
-
-const DebounceInput: React.FC<DebounceInputProps> = ({
-                                                       minChars,
-                                                       onDebounce,
-                                                       debounceTimeout,
-                                                       ...rest
-                                                     }) => {
-  const [inputValue, setInputValue] = React.useState('');
-  const timerRef = React.useRef<ReturnType<typeof setTimeout>>();
-
-  const handleChange = (event: React.ChangeEvent<HTMLInputElement>) => {
-    const newValue = event.target.value;
-    setInputValue(newValue); // Update internal state immediately for a responsive UI
-
-    if (timerRef.current) {
-      clearTimeout(timerRef.current);
-    }
-
-    timerRef.current = setTimeout(() => {
-      // Don't call onDebounce if the input value is too short
-      if (newValue && minChars && newValue?.length < minChars)
-        return;
-      onDebounce(newValue); // Call onDebounce after the debounce timeout
-    }, debounceTimeout);
-  };
-
-  React.useEffect(() => {
-    return () => {
-      if (timerRef.current) {
-        clearTimeout(timerRef.current);
-      }
-    };
-  }, []);
-
-  const handleClear = () => {
-    setInputValue(''); // Clear internal state
-    onDebounce(''); // Call onDebounce with empty string
-  };
-
-  return (
-    <Input
-      {...rest}
-      value={inputValue}
-      onChange={handleChange}
-      aria-label={rest['aria-label'] || 'Search'}
-      startDecorator={<SearchIcon />}
-      endDecorator={
-        <Box sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
-          {!!inputValue && (
-            <IconButton size='sm' aria-label='Clear search' onClick={handleClear}>
-              <ClearIcon sx={{ fontSize: 'xl' }} />
-            </IconButton>
-          )}
-          {rest.endDecorator}
-        </Box>
-      }
-    />
-  );
-};
-
-export const DebounceInputMemo = React.memo(DebounceInput);

src/common/components/DebouncedInput.tsx

@@ -0,0 +1,113 @@
+import * as React from 'react';
+
+import type { InputProps } from '@mui/joy/Input';
+import { Box, IconButton, Input } from '@mui/joy';
+import ClearIcon from '@mui/icons-material/Clear';
+import SearchIcon from '@mui/icons-material/Search';
+
+
+type DebounceInputProps = Omit<InputProps, 'onChange'> & {
+  /**
+   * When true, this will not give up the focus on the input field, and aggressively
+   * refocus it after the debounce (assuming the callee will cascade a removal, which
+   * is the case for Joy UI Select components).
+   */
+  aggressiveRefocus?: boolean;
+  debounceTimeout: number;
+  minChars?: number;
+  onDebounce: (value: string) => void;
+};
+
+const DebouncedInput: React.FC<DebounceInputProps> = (props: DebounceInputProps) => {
+
+  // state
+  const [inputValue, setInputValue] = React.useState('');
+  const inputRef = React.useRef<HTMLInputElement>(null);
+  const debounceTimerRef = React.useRef<ReturnType<typeof setTimeout>>();
+  const refocusTimerRef = React.useRef<ReturnType<typeof setTimeout>>();
+
+  // derived state
+  const { debounceTimeout, minChars, onDebounce, aggressiveRefocus, ...rest } = props;
+
+  // callbacks
+
+  const handleChange = React.useCallback((event: React.ChangeEvent<HTMLInputElement>) => {
+    const newValue = event.target.value;
+    setInputValue(newValue); // Update internal state immediately for a responsive UI
+
+    if (debounceTimerRef.current)
+      clearTimeout(debounceTimerRef.current);
+
+    debounceTimerRef.current = setTimeout(() => {
+      // reset the timer
+      debounceTimerRef.current = undefined;
+
+      // Don't call onDebounce if the input value is too short
+      if (newValue && minChars && newValue?.length < minChars)
+        return;
+
+      // Call onDebounce with the new value
+      onDebounce(newValue);
+
+      // If requested, get back the focus
+      if (aggressiveRefocus) {
+        if (refocusTimerRef.current)
+          clearTimeout(refocusTimerRef.current);
+
+        refocusTimerRef.current = setTimeout(() => {
+          refocusTimerRef.current = undefined;
+          inputRef.current?.focus();
+        }, 20);
+      }
+    }, debounceTimeout);
+  }, [debounceTimeout, aggressiveRefocus, minChars, onDebounce]);
+
+  const handleClear = React.useCallback(() => {
+    setInputValue(''); // Clear internal state
+    onDebounce(''); // Call onDebounce with empty string
+  }, [onDebounce]);
+
+
+  // Clear all timers on unmount
+  React.useEffect(() => {
+    return () => {
+      if (debounceTimerRef.current)
+        clearTimeout(debounceTimerRef.current);
+      if (refocusTimerRef.current)
+        clearTimeout(refocusTimerRef.current);
+    };
+  }, []);
+
+
+  return (
+    <Input
+      {...rest}
+      value={inputValue}
+      onChange={handleChange}
+      aria-label={rest['aria-label'] || 'Search'}
+      startDecorator={<SearchIcon />}
+      onKeyDownCapture={!aggressiveRefocus ? undefined : (event) => {
+        /* We stop the propagation of the event to prevent the parent component from handling it.
+         * This is useful only when used inside a Select with options, as the select is eager to capture
+         * the focus at every keystroke. This way we keep the focus.
+         */
+        event.stopPropagation();
+      }}
+      endDecorator={
+        <Box sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
+          {!!inputValue && (
+            <IconButton size='sm' aria-label='Clear search' onClick={handleClear}>
+              <ClearIcon sx={{ fontSize: 'xl' }} />
+            </IconButton>
+          )}
+          {rest.endDecorator}
+        </Box>
+      }
+      slotProps={!aggressiveRefocus ? undefined : {
+        input: { ref: inputRef },
+      }}
+    />
+  );
+};
+
+export const DebouncedInputMemo = React.memo(DebouncedInput);