fix: focus last pane correctness issues (#5344)

* fix: correctness issues with focus-last-pane

* rustfmt

* add pr

3 files changed, 57 insertions, 4 deletions

diff --git a/zellij-server/src/panes/floating_panes/mod.rs b/zellij-server/src/panes/floating_panes/mod.rs
index 5dbc7fb2..b9afb18e 100644
--- a/zellij-server/src/panes/floating_panes/mod.rs
+++ b/zellij-server/src/panes/floating_panes/mod.rs
@@ -741,11 +741,33 @@ impl FloatingPanes {
     }
 
     pub fn focus_last_pane(&mut self, client_id: ClientId) {
-        if let Some(pane_id) = self.active_panes.get_last(&client_id).copied() {
-            self.focus_pane(pane_id, client_id);
+        let currently_active_pane_id = self.active_panes.get(&client_id).copied();
+        let last_pane_id = self
+            .active_panes
+            .get_last(&client_id)
+            .copied()
+            .filter(|last_pane_id| {
+                Some(*last_pane_id) != currently_active_pane_id
+                    && self
+                        .panes
+                        .get(last_pane_id)
+                        .map(|p| p.selectable())
+                        .unwrap_or(false)
+            });
+        let target_pane_id = last_pane_id
+            .or_else(|| self.most_recently_active_selectable_pane(currently_active_pane_id));
+        if let Some(target_pane_id) = target_pane_id {
+            self.focus_pane(target_pane_id, client_id);
             self.set_force_render();
         }
     }
+    fn most_recently_active_selectable_pane(&self, exclude: Option<PaneId>) -> Option<PaneId> {
+        self.panes
+            .iter()
+            .filter(|(pane_id, pane)| Some(**pane_id) != exclude && pane.selectable())
+            .max_by_key(|(_pane_id, pane)| pane.active_at())
+            .map(|(pane_id, _pane)| *pane_id)
+    }
 
     pub fn move_active_pane_down(&mut self, client_id: ClientId) {
         if let Some(active_pane_id) = self.active_panes.get(&client_id) {
@@ -929,6 +951,11 @@ impl FloatingPanes {
         let next_active_pane_id = next_active_pane_candidates
             .last()
             .map(|(pane_id, _pane)| **pane_id);
+        let last_active_pane_id = next_active_pane_candidates
+            .iter()
+            .rev()
+            .nth(1)
+            .map(|(pane_id, _pane)| **pane_id);
 
         let connected_clients: HashSet<ClientId> =
             self.connected_clients.borrow().iter().copied().collect();
@@ -943,6 +970,10 @@ impl FloatingPanes {
                         self.active_panes
                             .insert(client_id, next_active_pane_id, &mut self.panes);
                         self.focus_pane(next_active_pane_id, client_id);
+                        if let Some(last_active_pane_id) = last_active_pane_id {
+                            self.active_panes
+                                .set_last_pane(client_id, last_active_pane_id);
+                        }
                     },
                     None => {
                         self.defocus_pane(client_id);
diff --git a/zellij-server/src/panes/tiled_panes/mod.rs b/zellij-server/src/panes/tiled_panes/mod.rs
index 0c17aa01..d9448641 100644
--- a/zellij-server/src/panes/tiled_panes/mod.rs
+++ b/zellij-server/src/panes/tiled_panes/mod.rs
@@ -1911,6 +1911,9 @@ impl TiledPanes {
             self.reapply_pane_frames();
         }
 
+        if active_pane_id != next_active_pane_id {
+            self.active_panes.set_last_pane(client_id, active_pane_id);
+        }
         self.active_panes
             .insert(client_id, next_active_pane_id, &mut self.panes);
         self.set_pane_active_at(next_active_pane_id);
@@ -1938,6 +1941,9 @@ impl TiledPanes {
                 .expand_pane(&next_active_pane_id);
             self.reapply_pane_frames();
         }
+        if active_pane_id != next_active_pane_id {
+            self.active_panes.set_last_pane(client_id, active_pane_id);
+        }
         self.active_panes
             .insert(client_id, next_active_pane_id, &mut self.panes);
         self.set_pane_active_at(next_active_pane_id);
@@ -1945,11 +1951,17 @@ impl TiledPanes {
     }
     pub fn focus_last_pane(&mut self, client_id: ClientId) {
         let Some(last_pane_id) = self.get_last_pane_id(client_id) else {
+            self.focus_previous_pane(client_id);
             return;
         };
 
-        let previously_active_pane_id = self.active_panes.get(&client_id).unwrap();
-        let previously_active_pane = self.panes.get_mut(previously_active_pane_id).unwrap();
+        let previously_active_pane_id = *self.active_panes.get(&client_id).unwrap();
+        if !self.panes.contains_key(&last_pane_id) || last_pane_id == previously_active_pane_id {
+            self.focus_previous_pane(client_id);
+            return;
+        }
+
+        let previously_active_pane = self.panes.get_mut(&previously_active_pane_id).unwrap();
 
         previously_active_pane.set_should_render(true);
         // we render the full viewport to remove any ui elements that might have been
@@ -2538,6 +2550,11 @@ impl TiledPanes {
         let next_active_pane_id = next_active_pane_candidates
             .last()
             .map(|(pane_id, _pane)| **pane_id);
+        let last_active_pane_id = next_active_pane_candidates
+            .iter()
+            .rev()
+            .nth(1)
+            .map(|(pane_id, _pane)| **pane_id);
 
         match next_active_pane_id {
             Some(next_active_pane_id) => {
@@ -2559,6 +2576,10 @@ impl TiledPanes {
                                 next_active_pane_id,
                                 &mut self.panes,
                             );
+                            if let Some(last_active_pane_id) = last_active_pane_id {
+                                self.active_panes
+                                    .set_last_pane(client_id, last_active_pane_id);
+                            }
                         } else {
                             self.active_panes.remove_silent(&client_id);
                         }
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index dc9c0409..d3bc154f 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -4611,6 +4611,7 @@ impl Tab {
         }
         if self.floating_panes.panes_are_visible() {
             self.floating_panes.focus_last_pane(client_id);
+            return;
         } else if self.tiled_panes.fullscreen_is_active() {
             self.switch_last_pane_fullscreen(client_id);
             return;