🤖 refactor: deflake ChatInput focus for Storybook

Author: ammar-agent

src/browser/components/ChatInput/index.tsx

@@ -325,6 +325,21 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
     }
   }, [variant, defaultModel, storageKeys.modelKey]);
 
+  // ChatInputSection is always present and is a convenient place to surface focus scheduling
+  // state for Storybook/tests without adding story-only props.
+  const chatInputSectionRef = useRef<HTMLDivElement | null>(null);
+  const handleChatInputSectionRef = useCallback((el: HTMLDivElement | null) => {
+    chatInputSectionRef.current = el;
+    // Set an initial value once (avoid overriding later transitions).
+    if (el && !el.hasAttribute("data-autofocus-state")) {
+      el.setAttribute("data-autofocus-state", "pending");
+    }
+  }, []);
+
+  const setChatInputAutoFocusState = useCallback((state: "pending" | "done") => {
+    chatInputSectionRef.current?.setAttribute("data-autofocus-state", state);
+  }, []);
+
   const focusMessageInput = useCallback(() => {
     const element = inputRef.current;
     if (!element || element.disabled) {
@@ -610,16 +625,50 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
   }, [voiceInput, setToast]);
 
   // Auto-focus chat input when workspace changes (workspace only)
+  //
+  // This is intentionally NOT a setTimeout-based delay. Fixed sleeps are prone to races
+  // (especially in Storybook) and can still fire after other UI interactions.
   const workspaceIdForFocus = variant === "workspace" ? props.workspaceId : null;
   useEffect(() => {
     if (variant !== "workspace") return;
 
-    // Small delay to ensure DOM is ready and other components have settled
-    const timer = setTimeout(() => {
+    const maxFrames = 10;
+    setChatInputAutoFocusState("pending");
+
+    let cancelled = false;
+    let rafId: number | null = null;
+    let attempts = 0;
+
+    const step = () => {
+      if (cancelled) return;
+
+      attempts += 1;
       focusMessageInput();
-    }, 100);
-    return () => clearTimeout(timer);
-  }, [variant, workspaceIdForFocus, focusMessageInput]);
+
+      const input = inputRef.current;
+      const isFocused = !!input && document.activeElement === input;
+      const isDone = isFocused || attempts >= maxFrames;
+
+      if (isDone) {
+        setChatInputAutoFocusState("done");
+        return;
+      }
+
+      rafId = requestAnimationFrame(step);
+    };
+
+    // Start on the next frame so the textarea is mounted and ready.
+    rafId = requestAnimationFrame(step);
+
+    return () => {
+      cancelled = true;
+      if (rafId !== null) {
+        cancelAnimationFrame(rafId);
+      }
+      // Ensure we never leave the attribute stuck in "pending".
+      setChatInputAutoFocusState("done");
+    };
+  }, [variant, workspaceIdForFocus, focusMessageInput, setChatInputAutoFocusState]);
 
   // Handle paste events to extract images
   const handlePaste = useCallback((e: React.ClipboardEvent<HTMLTextAreaElement>) => {
@@ -1371,6 +1420,7 @@ const ChatInputInner: React.FC<ChatInputProps> = (props) => {
 
       {/* Input section - centered card for creation, bottom bar for workspace */}
       <div
+        ref={handleChatInputSectionRef}
         className={cn(
           "relative flex flex-col gap-1",
           variant === "creation"

src/browser/stories/App.bash.stories.tsx

@@ -64,9 +64,19 @@ async function expandAllBashTools(canvasElement: HTMLElement) {
   }
 
   // Avoid leaving focus on a tool header.
-  // ChatInput also auto-focuses on a 100ms timer on mount/workspace changes; wait for
-  // that to fire before blurring so Storybook screenshots are deterministic.
-  await new Promise((resolve) => setTimeout(resolve, 150));
+  // Wait for ChatInput's auto-focus attempt to finish (no timing-based sleeps).
+  await waitFor(
+    () => {
+      const state = canvasElement
+        .querySelector('[data-component="ChatInputSection"]')
+        ?.getAttribute("data-autofocus-state");
+      if (state !== "done") {
+        throw new Error("ChatInput auto-focus not finished");
+      }
+    },
+    { timeout: 5000 }
+  );
+
   (document.activeElement as HTMLElement | null)?.blur?.();
 }
 

src/browser/stories/App.chat.stories.tsx

@@ -330,8 +330,18 @@ export const GenericTool: AppStory = {
       },
       { timeout: 5000 }
     );
-    // Wait for any auto-focus timers (ChatInput has 100ms delay), then blur
-    await new Promise((resolve) => setTimeout(resolve, 150));
+    // Wait for ChatInput's auto-focus attempt to finish (no timing-based sleeps), then blur
+    await waitFor(
+      () => {
+        const state = canvasElement
+          .querySelector('[data-component="ChatInputSection"]')
+          ?.getAttribute("data-autofocus-state");
+        if (state !== "done") {
+          throw new Error("ChatInput auto-focus not finished");
+        }
+      },
+      { timeout: 5000 }
+    );
     (document.activeElement as HTMLElement)?.blur();
   },
 };

src/browser/stories/App.reviews.stories.tsx

@@ -273,8 +273,18 @@ export const BulkReviewActions: AppStory = {
       await userEvent.click(bannerButton);
     });
 
-    // Wait for any auto-focus timers, then blur
-    await new Promise((resolve) => setTimeout(resolve, 150));
+    // Wait for ChatInput's auto-focus attempt to finish (no timing-based sleeps), then blur
+    await waitFor(
+      () => {
+        const state = canvasElement
+          .querySelector('[data-component="ChatInputSection"]')
+          ?.getAttribute("data-autofocus-state");
+        if (state !== "done") {
+          throw new Error("ChatInput auto-focus not finished");
+        }
+      },
+      { timeout: 5000 }
+    );
     (document.activeElement as HTMLElement)?.blur();
   },
 };