From 9194d7f6bb1ef132e32d8056296c90ce01bdfe05 Mon Sep 17 00:00:00 2001
From: Floyd Wang <gassnake999@gmail.com>
Date: Wed, 5 Mar 2025 16:41:43 +0800
Subject: [PATCH] popup_menu: Ensure focus is returned after menu is closed
 (#689)

---
 crates/story/src/lib.rs         |  3 +--
 crates/story/src/popup_story.rs |  3 +++
 crates/ui/src/popup_menu.rs     | 41 ++++++++-------------------------
 3 files changed, 14 insertions(+), 33 deletions(-)

diff --git a/crates/story/src/lib.rs b/crates/story/src/lib.rs
index c0b50345..8dd66e1d 100644
--- a/crates/story/src/lib.rs
+++ b/crates/story/src/lib.rs
@@ -483,8 +483,7 @@ impl Panel for StoryContainer {
     }
 
     fn popup_menu(&self, menu: PopupMenu, _window: &Window, _cx: &App) -> PopupMenu {
-        menu.track_focus(&self.focus_handle)
-            .menu("Info", Box::new(ShowPanelInfo))
+        menu.menu("Info", Box::new(ShowPanelInfo))
     }
 
     fn toolbar_buttons(&self, _window: &mut Window, _cx: &mut App) -> Option<Vec<Button>> {
diff --git a/crates/story/src/popup_story.rs b/crates/story/src/popup_story.rs
index ec9052b2..7df4e9e2 100644
--- a/crates/story/src/popup_story.rs
+++ b/crates/story/src/popup_story.rs
@@ -111,6 +111,9 @@ impl PopupStory {
 
     fn new(window: &mut Window, cx: &mut Context<Self>) -> Self {
         let form = Form::new(window, cx);
+
+        cx.focus_self(window);
+
         Self {
             form,
             focus_handle: cx.focus_handle(),
diff --git a/crates/ui/src/popup_menu.rs b/crates/ui/src/popup_menu.rs
index 83f23608..1d803a1a 100644
--- a/crates/ui/src/popup_menu.rs
+++ b/crates/ui/src/popup_menu.rs
@@ -107,7 +107,7 @@ pub struct PopupMenu {
     scroll_handle: ScrollHandle,
     scroll_state: Rc<Cell<ScrollbarState>>,
 
-    action_focus_handle: Option<FocusHandle>,
+    previous_focus_handle: Option<FocusHandle>,
     _subscriptions: Vec<Subscription>,
 }
 
@@ -128,7 +128,7 @@ impl PopupMenu {
 
             let menu = Self {
                 focus_handle,
-                action_focus_handle: None,
+                previous_focus_handle: window.focused(cx),
                 parent_menu: None,
                 menu_items: Vec::new(),
                 selected_index: None,
@@ -147,12 +147,6 @@ impl PopupMenu {
         })
     }
 
-    /// Bind the focus handle of the menu, when clicked, it will focus back to this handle and then dispatch the action
-    pub fn track_focus(mut self, focus_handle: &FocusHandle) -> Self {
-        self.action_focus_handle = Some(focus_handle.clone());
-        self
-    }
-
     /// Set min width of the popup menu, default is 120px
     pub fn min_w(mut self, width: impl Into<Pixels>) -> Self {
         self.min_width = width.into();
@@ -255,27 +249,7 @@ impl PopupMenu {
     }
 
     fn wrap_handler(&self, action: Box<dyn Action>) -> Rc<dyn Fn(&mut Window, &mut App)> {
-        let action_focus_handle = self.action_focus_handle.clone();
-
         Rc::new(move |window, cx| {
-            window.activate_window();
-
-            // Focus back to the user expected focus handle
-            // Then the actions listened on that focus handle can be received
-            //
-            // For example:
-            //
-            // TabPanel
-            //   |- PopupMenu
-            //   |- PanelContent (actions are listened here)
-            //
-            // The `PopupMenu` and `PanelContent` are at the same level in the TabPanel
-            // If the actions are listened on the `PanelContent`,
-            // it can't receive the actions from the `PopupMenu`, unless we focus on `PanelContent`.
-            if let Some(handle) = action_focus_handle.as_ref() {
-                window.focus(&handle);
-            }
-
             window.dispatch_action(action.boxed_clone(), cx);
         })
     }
@@ -431,13 +405,18 @@ impl PopupMenu {
         }
     }
 
-    fn dismiss(&mut self, _: &Dismiss, _window: &mut Window, cx: &mut Context<Self>) {
+    fn dismiss(&mut self, _: &Dismiss, window: &mut Window, cx: &mut Context<Self>) {
         if self.active_submenu().is_some() {
             return;
         }
 
         cx.emit(DismissEvent);
 
+        // Focus back to the previous focused handle.
+        if let Some(previous_focus_handle) = self.previous_focus_handle.as_ref() {
+            window.focus(previous_focus_handle);
+        }
+
         let Some(parent_menu) = self.parent_menu.clone() else {
             return;
         };
@@ -445,7 +424,7 @@ impl PopupMenu {
         // Dismiss parent menu, when this menu is dismissed
         _ = parent_menu.update(cx, |view, cx| {
             view.hovered_menu_ix = None;
-            view.dismiss(&Dismiss, _window, cx);
+            view.dismiss(&Dismiss, window, cx);
         });
     }
 
@@ -528,7 +507,7 @@ impl Render for PopupMenu {
         v_flex()
             .id("popup-menu")
             .key_context("PopupMenu")
-            // .track_focus(&self.focus_handle)
+            .track_focus(&self.focus_handle)
             .on_action(cx.listener(Self::select_next))
             .on_action(cx.listener(Self::select_prev))
             .on_action(cx.listener(Self::confirm))