diff --git a/app/client/src/utils/hooks/useWidgetFocus/useWidgetFocus.tsx b/app/client/src/utils/hooks/useWidgetFocus/useWidgetFocus.tsx
index 50e2ee3a96..bc3772cda0 100644
--- a/app/client/src/utils/hooks/useWidgetFocus/useWidgetFocus.tsx
+++ b/app/client/src/utils/hooks/useWidgetFocus/useWidgetFocus.tsx
@@ -1,52 +1,70 @@
+import { useCallback, useRef } from "react";
 import { useSelector } from "react-redux";
-import { useCallback, useEffect, useRef } from "react";
 
+import { getIsAutoLayout } from "selectors/canvasSelectors";
 import { handleTab } from "./handleTab";
 import { CANVAS_WIDGET } from "./tabbable";
-import { getIsAutoLayout } from "selectors/canvasSelectors";
 
 function useWidgetFocus(): (instance: HTMLElement | null) => void {
   const ref = useRef<HTMLElement | null>();
   const isAutoLayout = useSelector(getIsAutoLayout);
 
-  // This is a callback that will be called when the ref is set
-  const setRef = useCallback((node: HTMLElement | null) => {
-    if (node === null) return;
-
-    if (ref.current === node) return;
-
-    ref.current = node;
-
-    return ref;
-  }, []);
-
-  useEffect(() => {
-    if (isAutoLayout) return;
-
-    if (!ref.current) return;
-
-    const handleKeyDown = (event: KeyboardEvent) => {
-      if (event.key === "Tab") handleTab(event);
-    };
-
-    // TODO: Fix this the next time the file is edited
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const handleClick = (event: any) => {
-      const target = event.target as HTMLElement;
-
-      if (target.matches(CANVAS_WIDGET)) {
-        target.focus();
+  const attachEventListeners = useCallback(
+    (element: HTMLElement) => {
+      if (isAutoLayout) {
+        return () => {}; // Return empty cleanup function
       }
-    };
 
-    ref.current.addEventListener("keydown", handleKeyDown);
-    ref.current.addEventListener("click", handleClick);
+      const handleKeyDown = (event: KeyboardEvent) => {
+        if (event.key === "Tab") {
+          handleTab(event);
+        }
+      };
 
-    return () => {
-      ref?.current && ref.current.removeEventListener("keydown", handleKeyDown);
-      ref?.current && ref.current.removeEventListener("click", handleClick);
-    };
-  }, []);
+      // TODO: Fix this the next time the file is edited
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const handleClick = (event: any) => {
+        const target = event.target as HTMLElement;
+
+        if (target.matches(CANVAS_WIDGET)) {
+          target.focus();
+        }
+      };
+
+      element.addEventListener("keydown", handleKeyDown);
+      element.addEventListener("click", handleClick);
+
+      // Return cleanup function
+      return () => {
+        element.removeEventListener("keydown", handleKeyDown);
+        element.removeEventListener("click", handleClick);
+      };
+    },
+    [isAutoLayout],
+  );
+
+  // This is a callback that will be called when the ref is set
+  const setRef = useCallback(
+    (node: HTMLElement | null) => {
+      if (node === null) {
+        ref.current = null;
+
+        return;
+      }
+
+      if (ref.current === node) {
+        return;
+      }
+
+      ref.current = node;
+
+      // Attach event listeners immediately when ref is set
+      attachEventListeners(node);
+
+      return ref;
+    },
+    [attachEventListeners],
+  );
 
   return setRef;
 }