🎨 Palette: Prevent Focus Loss on Dynamic Component Unmount

src/components/chat/client.test.tsx

@@ -111,6 +111,27 @@ describe("ChatWidgetClient", () => {
     );
   });
 
+  it("should focus the input textarea after a suggested question is triggered", () => {
+    mockUseChatWidgetResult.isOpen = true;
+    mockUseChatWidgetResult.messages = [];
+
+    const { getByText, getByLabelText } = render(<ChatWidgetClient />);
+    const textarea = getByLabelText("Ask a question to Miro");
+    const focusSpy = vi.spyOn(textarea, "focus");
+
+    mockHandleSuggestedQuestion.mockImplementation(() => {
+      textarea.focus();
+    });
+
+    const suggestedQuestion = getByText("What was his PhD research about?");
+
+    act(() => {
+      suggestedQuestion.click();
+    });
+
+    expect(focusSpy).toHaveBeenCalled();
+  });
+
   it("should render messages list, typing indicator and stop button when loading", () => {
     mockUseChatWidgetResult.isOpen = true;
     mockUseChatWidgetResult.isLoading = true;

src/components/chat/use-chat-widget.ts

@@ -102,6 +102,9 @@ export function useChatWidget() {
       }
       try {
         await sendMessage({ text: currentInput });
+        if (inputRef.current) {
+          inputRef.current.focus();
+        }
       } catch (err) {
         console.error("Failed to send message:", err);
       }
@@ -114,6 +117,9 @@ export function useChatWidget() {
       if (isLoading) return;
       try {
         await sendMessage({ text: question });
+        if (inputRef.current) {
+          inputRef.current.focus();
+        }
       } catch (err) {
         console.error("Failed to send suggested question:", err);
       }

src/components/components.test.tsx

@@ -109,15 +109,14 @@ describe("UI Components", () => {
   });
 
   describe("ScrollToTopButton (upArrow)", () => {
-    it("should show/hide on scroll and trigger window.scrollTo on click", () => {
-      const { container } = render(<ScrollToTopButton />);
+    it("should show/hide on scroll and trigger window.scrollTo on click, shifting focus to main-content or body", () => {
+      // Create a dummy main-content div in the test DOM
+      const mainContent = document.createElement("div");
+      mainContent.id = "main-content";
+      const focusSpy = vi.spyOn(mainContent, "focus");
+      document.body.appendChild(mainContent);
 
-      // Initially not visible (window.scrollY is 0)
-      window.scrollY = 0;
-      act(() => {
-        window.dispatchEvent(new Event("scroll"));
-      });
-      expect(container.querySelector(".scroll-button")).not.toBeInTheDocument();
+      const { container } = render(<ScrollToTopButton />);
 
       // Scroll down
       window.scrollY = 400;
@@ -137,6 +136,31 @@ describe("UI Components", () => {
         top: 0,
         behavior: "smooth",
       });
+
+      // Verify that focus was shifted to main-content
+      expect(focusSpy).toHaveBeenCalledWith({ preventScroll: true });
+
+      // Clean up
+      document.body.removeChild(mainContent);
+    });
+
+    it("should fallback to document.body.focus if main-content is not found", () => {
+      const bodyFocusSpy = vi.spyOn(document.body, "focus");
+
+      const { container } = render(<ScrollToTopButton />);
+
+      // Scroll down
+      window.scrollY = 400;
+      act(() => {
+        window.dispatchEvent(new Event("scroll"));
+      });
+
+      const btn = container.querySelector(".scroll-button");
+      act(() => {
+        (btn as HTMLButtonElement).click();
+      });
+
+      expect(bodyFocusSpy).toHaveBeenCalledWith({ preventScroll: true });
     });
   });
 

src/components/upArrow.tsx

@@ -30,6 +30,12 @@ const ScrollToTopButton = () => {
       top: 0,
       behavior: "smooth",
     });
+    const mainContent = document.getElementById("main-content");
+    if (mainContent) {
+      mainContent.focus({ preventScroll: true });
+    } else {
+      document.body.focus({ preventScroll: true });
+    }
   };
 
   return (