fix: focus last pane correctness issues (#5344)

* fix: correctness issues with focus-last-pane

* rustfmt

* add pr

6 files changed, 251 insertions, 5 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 404b0c93..8a568431 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 * feat: PWA support for the web client (manifest + icons + iOS meta tags) so the page can be installed as a standalone app (https://github.com/zellij-org/zellij/pull/5184)
 * feat: mobile UI (https://github.com/zellij-org/zellij/pull/5241)
 * chore: new integration testing infra (https://github.com/zellij-org/zellij/pull/5269)
-* feat: FocusLastPane bindable action (https://github.com/zellij-org/zellij/pull/4241)
+* feat: FocusLastPane bindable action (https://github.com/zellij-org/zellij/pull/4241 and https://github.com/zellij-org/zellij/pull/5344)
 * fix: live reload on theme dir changes (https://github.com/zellij-org/zellij/pull/5135)
 * feat: new title-frames UI (https://github.com/zellij-org/zellij/pull/5318)
 * chore: upgrade wasmi to 1.1.0 (https://github.com/zellij-org/zellij/pull/5319)
diff --git a/zellij-integration-tests/tests/focus_last_pane.rs b/zellij-integration-tests/tests/focus_last_pane.rs
new file mode 100644
index 00000000..9058c8b4
--- /dev/null
+++ b/zellij-integration-tests/tests/focus_last_pane.rs
@@ -0,0 +1,165 @@
+#![cfg(unix)]
+
+use insta::assert_snapshot;
+use zellij_integration_tests::{
+    claim_first_terminal_and_wait_for_prompt, col, keys, normalized,
+    split_down_and_wait_for_prompt, split_right_and_wait_for_prompt, FakePtyHandle, Size,
+    TestRunner, TestSession, PROMPT, TERMINAL_SIZE,
+};
+
+const FOCUS_KEYS_CONFIG: &str = r#"
+keybinds {
+    normal {
+        bind "Alt m" { FocusLastPane; }
+        bind "Alt ." { FocusNextPane; }
+    }
+}
+"#;
+
+fn start_with_focus_keys(size: Size) -> TestSession {
+    TestRunner::new(size).with_config(FOCUS_KEYS_CONFIG).start()
+}
+
+fn split_focused_pane_down(zellij: &TestSession) -> FakePtyHandle {
+    zellij.send_stdin(&keys::ctrl('p'));
+    zellij.send_stdin(&keys::key('d'));
+    let terminal = zellij.expect_pty_spawn();
+    terminal.output(PROMPT);
+    zellij.wait_until(
+        "newly split lower pane focused below the middle pane",
+        |grid| grid.status_bar_appears() && grid.cursor.map_or(false, |cursor| cursor.y > 13),
+    );
+    terminal
+}
+
+#[test]
+fn focus_last_pane_toggles_focus_back_to_the_previous_pane() {
+    let mut zellij = start_with_focus_keys(TERMINAL_SIZE);
+    claim_first_terminal_and_wait_for_prompt(&zellij);
+    split_right_and_wait_for_prompt(&zellij);
+
+    zellij.send_stdin(&keys::alt('m'));
+    let grid = zellij.wait_until(
+        "focus_last_pane moved focus back to the left pane",
+        |grid| grid.status_bar_appears() && grid.cursor_is_at(col(2).row(2)),
+    );
+    assert_snapshot!(normalized(&grid));
+    zellij.quit();
+}
+
+#[test]
+fn focus_last_pane_returns_to_earlier_pane_after_closing_focused_pane() {
+    let mut zellij = start_with_focus_keys(TERMINAL_SIZE);
+    claim_first_terminal_and_wait_for_prompt(&zellij);
+    split_down_and_wait_for_prompt(&zellij);
+    let bottom_terminal = split_focused_pane_down(&zellij);
+
+    bottom_terminal.exit(Some(0));
+    zellij.wait_until(
+        "closing the bottom pane returns focus to the middle pane",
+        |grid| grid.status_bar_appears() && grid.cursor_is_at(col(2).row(13)),
+    );
+
+    zellij.send_stdin(&keys::alt('m'));
+    zellij.wait_until(
+        "focus_last_pane moves focus to the pane active before the closed one",
+        |grid| grid.status_bar_appears() && grid.cursor_is_at(col(2).row(2)),
+    );
+    zellij.quit();
+}
+
+#[test]
+fn focus_last_pane_refills_after_cycling_focus() {
+    let mut zellij = start_with_focus_keys(TERMINAL_SIZE);
+    claim_first_terminal_and_wait_for_prompt(&zellij);
+    split_right_and_wait_for_prompt(&zellij);
+
+    zellij.send_stdin(&keys::alt('.'));
+    zellij.wait_until("focus cycled away from the right pane", |grid| {
+        grid.status_bar_appears() && grid.cursor.map_or(false, |cursor| cursor.x < 60)
+    });
+
+    zellij.send_stdin(&keys::alt('m'));
+    zellij.wait_until(
+        "focus_last_pane returns to the pane focused before cycling",
+        |grid| grid.status_bar_appears() && grid.cursor_is_at(col(62).row(2)),
+    );
+    zellij.quit();
+}
+
+fn open_floating_pane_below(
+    zellij: &TestSession,
+    previous_cursor_row: usize,
+) -> (FakePtyHandle, usize) {
+    zellij.send_stdin(&keys::alt('n'));
+    let terminal = zellij.expect_pty_spawn();
+    terminal.output(PROMPT);
+    let grid = zellij.wait_until("new floating pane focused below the previous one", |grid| {
+        grid.cursor
+            .map_or(false, |cursor| cursor.y > previous_cursor_row)
+    });
+    (terminal, grid.cursor.unwrap().y)
+}
+
+#[test]
+fn focus_last_floating_pane_returns_to_earlier_pane_after_closing_focused_pane() {
+    let mut zellij = start_with_focus_keys(TERMINAL_SIZE);
+    claim_first_terminal_and_wait_for_prompt(&zellij);
+
+    zellij.send_stdin(&keys::ctrl('p'));
+    zellij.send_stdin(&keys::key('w'));
+    let first_floating = zellij.expect_pty_spawn();
+    first_floating.output(PROMPT);
+    let first_floating_cursor = zellij
+        .wait_until("first floating pane focused", |grid| {
+            grid.contains("Pane #2") && grid.cursor.is_some()
+        })
+        .cursor
+        .unwrap();
+
+    let (_second_floating, second_floating_row) =
+        open_floating_pane_below(&zellij, first_floating_cursor.y);
+    let (third_floating, _third_floating_row) =
+        open_floating_pane_below(&zellij, second_floating_row);
+
+    third_floating.exit(Some(0));
+    zellij.wait_until(
+        "closing the focused floating pane returns focus to the previous floating pane",
+        |grid| {
+            grid.cursor
+                .map_or(false, |cursor| cursor.y == second_floating_row)
+        },
+    );
+
+    zellij.send_stdin(&keys::alt('m'));
+    zellij.wait_until(
+        "focus_last_pane moves focus to the floating pane active before the closed one",
+        |grid| grid.cursor == Some(first_floating_cursor),
+    );
+    zellij.quit();
+}
+
+#[test]
+fn focus_last_pane_with_floating_visible_leaves_tiled_focus_untouched() {
+    let mut zellij = start_with_focus_keys(TERMINAL_SIZE);
+    claim_first_terminal_and_wait_for_prompt(&zellij);
+    split_right_and_wait_for_prompt(&zellij);
+
+    zellij.send_stdin(&keys::ctrl('p'));
+    zellij.send_stdin(&keys::key('w'));
+    let floating_terminal = zellij.expect_pty_spawn();
+    floating_terminal.output(PROMPT);
+    zellij.wait_until("floating pane focused over the tiled panes", |grid| {
+        grid.contains("Pane #3") && grid.cursor.is_some()
+    });
+
+    zellij.send_stdin(&keys::alt('m'));
+
+    zellij.send_stdin(&keys::ctrl('p'));
+    zellij.send_stdin(&keys::key('w'));
+    zellij.wait_until(
+        "the tiled pane focused before opening floating panes stays focused",
+        |grid| grid.status_bar_appears() && grid.cursor_is_at(col(62).row(2)),
+    );
+    zellij.quit();
+}
diff --git a/zellij-integration-tests/tests/snapshots/focus_last_pane__focus_last_pane_toggles_focus_back_to_the_previous_pane.snap b/zellij-integration-tests/tests/snapshots/focus_last_pane__focus_last_pane_toggles_focus_back_to_the_previous_pane.snap
new file mode 100644
index 00000000..bc52074f
--- /dev/null
+++ b/zellij-integration-tests/tests/snapshots/focus_last_pane__focus_last_pane_toggles_focus_back_to_the_previous_pane.snap
@@ -0,0 +1,28 @@
+---
+source: zellij-integration-tests/tests/focus_last_pane.rs
+expression: normalized(&grid)
+---
+ Zellij (test)  Tab #1  + 
+                           Pane #1                         │                           Pane #2
+$ █                                                        │$
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+                                                           │
+ Ctrl + <g> LOCK  <p> PANE  <t> TAB  <n> RESIZE  <h> MOVE  <s> SEARCH  <o> SESSION  <q> QUIT        ...
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;