[PAI-963] feat: enhance CustomSelect component with context for dropdown management (#8202)

* feat: enhance CustomSelect component with context for dropdown management

* refactor: streamline CustomSelect component structure and improve dropdown options rendering

packages/ui/src/dropdowns/custom-select.tsx

@@ -1,6 +1,6 @@
 import { Combobox } from "@headlessui/react";
 import { Check } from "lucide-react";
-import React, { useRef, useState } from "react";
+import React, { createContext, useCallback, useContext, useRef, useState } from "react";
 import { createPortal } from "react-dom";
 import { usePopper } from "react-popper";
 import { useOutsideClickDetector } from "@plane/hooks";
@@ -13,6 +13,9 @@ import { cn } from "../utils";
 // types
 import type { ICustomSelectItemProps, ICustomSelectProps } from "./helper";
 
+// Context to share the close handler with option components
+const DropdownContext = createContext<() => void>(() => {});
+
 function CustomSelect(props: ICustomSelectProps) {
   const {
     customButtonClassName = "",
@@ -42,26 +45,31 @@ function CustomSelect(props: ICustomSelectProps) {
     placement: placement ?? "bottom-start",
   });
 
-  const openDropdown = () => {
+  const openDropdown = useCallback(() => {
     setIsOpen(true);
     if (referenceElement) referenceElement.focus();
-  };
+  }, [referenceElement]);
-  const closeDropdown = () => setIsOpen(false);
+
+  const closeDropdown = useCallback(() => setIsOpen(false), []);
   const handleKeyDown = useDropdownKeyDown(openDropdown, closeDropdown, isOpen);
   useOutsideClickDetector(dropdownRef, closeDropdown);
 
-  const toggleDropdown = () => {
+  const toggleDropdown = useCallback(() => {
     if (isOpen) closeDropdown();
     else openDropdown();
-  };
+  }, [closeDropdown, isOpen, openDropdown]);
 
   return (
+    <DropdownContext.Provider value={closeDropdown}>
       <Combobox
         as="div"
         ref={dropdownRef}
         tabIndex={tabIndex}
         value={value}
-      onChange={onChange}
+        onChange={(val) => {
+          onChange?.(val);
+          closeDropdown();
+        }}
         className={cn("relative flex-shrink-0 text-left", className)}
         onKeyDown={handleKeyDown}
         disabled={disabled}
@@ -72,7 +80,7 @@ function CustomSelect(props: ICustomSelectProps) {
               <button
                 ref={setReferenceElement}
                 type="button"
-              className={`flex items-center justify-between gap-1 text-xs ${
+                className={`flex items-center justify-between gap-1 text-xs rounded ${
                   disabled ? "cursor-not-allowed text-custom-text-200" : "cursor-pointer hover:bg-custom-background-80"
                 } ${customButtonClassName}`}
                 onClick={toggleDropdown}
@@ -105,7 +113,7 @@ function CustomSelect(props: ICustomSelectProps) {
         </>
         {isOpen &&
           createPortal(
-          <Combobox.Options data-prevent-outside-click static>
+            <Combobox.Options data-prevent-outside-click>
               <div
                 className={cn(
                   "my-1 overflow-y-scroll rounded-md border-[0.5px] border-custom-border-300 bg-custom-background-100 px-2 py-2.5 text-xs shadow-custom-shadow-rg focus:outline-none min-w-48 whitespace-nowrap z-30",
@@ -130,11 +138,19 @@ function CustomSelect(props: ICustomSelectProps) {
             document.body
           )}
       </Combobox>
+    </DropdownContext.Provider>
   );
 }
 
 function Option(props: ICustomSelectItemProps) {
   const { children, value, className } = props;
+  const closeDropdown = useContext(DropdownContext);
+
+  const handleMouseDown = useCallback(() => {
+    // Close dropdown for both new and already-selected options.
+    requestAnimationFrame(() => closeDropdown());
+  }, [closeDropdown]);
+
   return (
     <Combobox.Option
       value={value}
@@ -149,10 +165,10 @@ function Option(props: ICustomSelectItemProps) {
       }
     >
       {({ selected }) => (
-        <>
+        <div onMouseDown={handleMouseDown} className="flex items-center justify-between gap-2 w-full">
           {children}
           {selected && <Check className="h-3.5 w-3.5 flex-shrink-0" />}
-        </>
+        </div>
       )}
     </Combobox.Option>
   );