Welcome screen del fix (#5160)

* fix: session-manager delete

* rustfmt

* add pr

4 files changed, 175 insertions, 40 deletions

diff --git a/default-plugins/session-manager/src/main.rs b/default-plugins/session-manager/src/main.rs
index 2c6e366c..6563a226 100644
--- a/default-plugins/session-manager/src/main.rs
+++ b/default-plugins/session-manager/src/main.rs
@@ -8,7 +8,7 @@ use uuid::Uuid;
 use zellij_tile::prelude::*;
 
 use new_session_info::NewSessionInfo;
-use single_screen::{SingleScreenMode, SingleScreenState, UnifiedSearchResult};
+use single_screen::{DeleteTarget, SingleScreenMode, SingleScreenState, UnifiedSearchResult};
 use ui::{
     components::{
         render_controls_line, render_error, render_new_session_block, render_prompt,
@@ -523,11 +523,26 @@ impl State {
             match key.bare_key {
                 BareKey::Char('y') if key.has_no_modifiers() => {
                     let all_other_sessions = self.sessions.all_other_sessions();
-                    kill_sessions(&all_other_sessions);
-                    self.reset_selected_index();
-                    self.search_term.clear();
-                    self.sessions
-                        .update_search_term(&self.search_term, &self.colors);
+                    let was_searching = self.sessions.is_searching;
+                    let prev_search_idx = self.sessions.selected_search_index;
+                    let prev_top_idx = self.sessions.selected_index.0;
+                    match kill_sessions(&all_other_sessions) {
+                        Ok(()) => {
+                            self.sessions
+                                .session_ui_infos
+                                .retain(|s| !all_other_sessions.contains(&s.name));
+                            self.sessions
+                                .update_search_term(&self.search_term, &self.colors);
+                            self.sessions.restore_selection_after_delete(
+                                was_searching,
+                                prev_search_idx,
+                                prev_top_idx,
+                            );
+                        },
+                        Err(e) => {
+                            self.show_error(&format!("Failed to kill sessions: {}", e));
+                        },
+                    }
                     self.show_kill_all_sessions_warning = false;
                     should_render = true;
                 },
@@ -599,11 +614,26 @@ impl State {
                 },
                 BareKey::Delete if key.has_no_modifiers() => {
                     if let Some(selected_session_name) = self.sessions.get_selected_session_name() {
-                        kill_sessions(&[selected_session_name]);
-                        self.reset_selected_index();
-                        self.search_term.clear();
-                        self.sessions
-                            .update_search_term(&self.search_term, &self.colors);
+                        let was_searching = self.sessions.is_searching;
+                        let prev_search_idx = self.sessions.selected_search_index;
+                        let prev_top_idx = self.sessions.selected_index.0;
+                        match kill_sessions(&[selected_session_name.clone()]) {
+                            Ok(()) => {
+                                self.sessions
+                                    .session_ui_infos
+                                    .retain(|s| s.name != selected_session_name);
+                                self.sessions
+                                    .update_search_term(&self.search_term, &self.colors);
+                                self.sessions.restore_selection_after_delete(
+                                    was_searching,
+                                    prev_search_idx,
+                                    prev_top_idx,
+                                );
+                            },
+                            Err(e) => {
+                                self.show_error(&format!("Failed to kill session: {}", e));
+                            },
+                        }
                     } else {
                         self.show_error("Must select session before killing it.");
                     }
@@ -725,7 +755,23 @@ impl State {
             match key.bare_key {
                 BareKey::Char('y') if key.has_no_modifiers() => {
                     let all_other_sessions = self.sessions.all_other_sessions();
-                    kill_sessions(&all_other_sessions);
+                    let previous_index = self.single_screen_state.selected_index;
+                    match kill_sessions(&all_other_sessions) {
+                        Ok(()) => {
+                            self.sessions
+                                .session_ui_infos
+                                .retain(|s| !all_other_sessions.contains(&s.name));
+                            self.single_screen_state.update_search_term(
+                                &self.sessions.session_ui_infos,
+                                &self.resurrectable_sessions.all_resurrectable_sessions,
+                            );
+                            self.single_screen_state
+                                .restore_selection_after_delete(previous_index);
+                        },
+                        Err(e) => {
+                            self.show_error(&format!("Failed to kill sessions: {}", e));
+                        },
+                    }
                     self.show_kill_all_sessions_warning = false;
                     should_render = true;
                 },
@@ -821,16 +867,35 @@ impl State {
                 should_render = true;
             },
             BareKey::Delete if key.has_no_modifiers() => {
-                if let Some(result) = self.single_screen_state.get_selected_result() {
-                    match result {
-                        UnifiedSearchResult::ActiveSession { session_name, .. } => {
-                            kill_sessions(&[session_name.clone()]);
+                let selected = self
+                    .single_screen_state
+                    .get_selected_result()
+                    .map(|r| r.as_delete_target());
+                if let Some(target) = selected {
+                    let previous_index = self.single_screen_state.selected_index;
+                    let outcome: Result<(), String> = match &target {
+                        DeleteTarget::Active(name) => kill_sessions(&[name.clone()]).map(|()| {
+                            self.sessions.session_ui_infos.retain(|s| s.name != *name);
+                        }),
+                        DeleteTarget::Resurrectable(name) => delete_dead_session(name).map(|()| {
+                            self.resurrectable_sessions
+                                .all_resurrectable_sessions
+                                .retain(|(n, _)| n != name);
+                        }),
+                    };
+                    match outcome {
+                        Ok(()) => {
+                            self.single_screen_state.update_search_term(
+                                &self.sessions.session_ui_infos,
+                                &self.resurrectable_sessions.all_resurrectable_sessions,
+                            );
+                            self.single_screen_state
+                                .restore_selection_after_delete(previous_index);
                         },
-                        UnifiedSearchResult::ResurrectableSession { session_name, .. } => {
-                            delete_dead_session(session_name);
+                        Err(e) => {
+                            self.show_error(&format!("Failed to delete session: {}", e));
                         },
                     }
-                    self.single_screen_state.selected_index = None;
                 }
                 should_render = true;
             },
diff --git a/default-plugins/session-manager/src/resurrectable_sessions.rs b/default-plugins/session-manager/src/resurrectable_sessions.rs
index a25d8b36..a2e42aee 100644
--- a/default-plugins/session-manager/src/resurrectable_sessions.rs
+++ b/default-plugins/session-manager/src/resurrectable_sessions.rs
@@ -260,34 +260,33 @@ impl ResurrectableSessions {
     }
     pub fn delete_selected_session(&mut self) {
         if self.is_searching {
-            self.selected_search_index
+            if let Some(search_result) = self
+                .selected_search_index
                 .and_then(|i| self.search_results.get(i))
-                .map(|search_result| delete_dead_session(&search_result.session_name));
-        } else {
-            self.selected_index
-                .and_then(|i| {
-                    if self.all_resurrectable_sessions.len() > i {
-                        // optimistic update
-                        if i == 0 {
-                            self.selected_index = None;
-                        } else if i == self.all_resurrectable_sessions.len().saturating_sub(1) {
-                            self.selected_index = Some(i.saturating_sub(1));
-                        }
-                        Some(self.all_resurrectable_sessions.remove(i))
-                    } else {
-                        None
-                    }
-                })
-                .map(|session_name_and_creation_time| {
-                    delete_dead_session(&session_name_and_creation_time.0)
-                });
+            {
+                let _ = delete_dead_session(&search_result.session_name);
+            }
+        } else if let Some(session_name_and_creation_time) = self.selected_index.and_then(|i| {
+            if self.all_resurrectable_sessions.len() > i {
+                // optimistic update
+                if i == 0 {
+                    self.selected_index = None;
+                } else if i == self.all_resurrectable_sessions.len().saturating_sub(1) {
+                    self.selected_index = Some(i.saturating_sub(1));
+                }
+                Some(self.all_resurrectable_sessions.remove(i))
+            } else {
+                None
+            }
+        }) {
+            let _ = delete_dead_session(&session_name_and_creation_time.0);
         }
     }
     fn delete_all_sessions(&mut self) {
         // optimistic update
         self.all_resurrectable_sessions = vec![];
         self.delete_all_dead_sessions_warning = false;
-        delete_all_dead_sessions();
+        let _ = delete_all_dead_sessions();
     }
     pub fn show_delete_all_sessions_warning(&mut self) {
         self.delete_all_dead_sessions_warning = true;
diff --git a/default-plugins/session-manager/src/session_list.rs b/default-plugins/session-manager/src/session_list.rs
index 7a66aabb..7d240f00 100644
--- a/default-plugins/session-manager/src/session_list.rs
+++ b/default-plugins/session-manager/src/session_list.rs
@@ -321,6 +321,31 @@ impl SessionList {
     pub fn reset_selected_index(&mut self) {
         self.selected_index.reset();
     }
+    /// After deleting one or more entries (and re-running `update_search_term`
+    /// to rebuild `search_results`), put the cursor back on a sensible
+    /// neighbour at the same numeric row -- the entry that took the deleted
+    /// row's slot, or the last entry if the deleted row was at the end. The
+    /// caller passes the indices captured **before** the deletion so this
+    /// method can clamp them to the new list/search lengths.
+    pub fn restore_selection_after_delete(
+        &mut self,
+        was_searching: bool,
+        prev_search_idx: Option<usize>,
+        prev_top_idx: Option<usize>,
+    ) {
+        if was_searching {
+            let len = self.search_results.len();
+            self.selected_search_index = clamp_index_after_delete(prev_search_idx, len);
+        } else {
+            // Tab / pane subselectors point into a session that has just
+            // shifted in the list; clear them so the top-level index alone
+            // describes the new selection.
+            self.selected_index.1 = None;
+            self.selected_index.2 = None;
+            self.selected_index.0 =
+                clamp_index_after_delete(prev_top_idx, self.session_ui_infos.len());
+        }
+    }
     pub fn has_session(&self, session_name: &str) -> bool {
         self.session_ui_infos.iter().any(|s| s.name == session_name)
     }
@@ -349,6 +374,17 @@ impl SessionList {
     }
 }
 
+/// Clamp a pre-delete index to a post-delete length. Returns the same
+/// numeric row when it still exists (so the cursor lands on whatever entry
+/// took the deleted row's slot), the last index when the deletion happened
+/// at the tail, or `None` when the list is now empty.
+pub fn clamp_index_after_delete(prev_index: Option<usize>, new_len: usize) -> Option<usize> {
+    if new_len == 0 {
+        return None;
+    }
+    Some(prev_index.unwrap_or(0).min(new_len - 1))
+}
+
 #[derive(Debug, Clone, Default)]
 pub struct SelectedIndex(pub Option<usize>, pub Option<usize>, pub Option<usize>);
 
diff --git a/default-plugins/session-manager/src/single_screen.rs b/default-plugins/session-manager/src/single_screen.rs
index fb7b0385..89ac91f7 100644
--- a/default-plugins/session-manager/src/single_screen.rs
+++ b/default-plugins/session-manager/src/single_screen.rs
@@ -39,7 +39,27 @@ pub enum UnifiedSearchResult {
     },
 }
 
+/// Which kind of entry was under the cursor when the user pressed `Delete`.
+/// Carries the session name so the host code can issue the right shim call
+/// (`kill_sessions` vs `delete_dead_session`) without re-borrowing the
+/// underlying `UnifiedSearchResult`.
+#[derive(Debug, Clone)]
+pub enum DeleteTarget {
+    Active(String),
+    Resurrectable(String),
+}
+
 impl UnifiedSearchResult {
+    pub fn as_delete_target(&self) -> DeleteTarget {
+        match self {
+            UnifiedSearchResult::ActiveSession { session_name, .. } => {
+                DeleteTarget::Active(session_name.clone())
+            },
+            UnifiedSearchResult::ResurrectableSession { session_name, .. } => {
+                DeleteTarget::Resurrectable(session_name.clone())
+            },
+        }
+    }
     pub fn session_name(&self) -> &str {
         match self {
             UnifiedSearchResult::ActiveSession { session_name, .. } => session_name.as_str(),
@@ -151,6 +171,21 @@ impl SingleScreenState {
         }
     }
 
+    /// After deleting an entry, ensure the cursor stays on a sensible row.
+    /// `update_search_term` already keeps the selection on the same entry by
+    /// name if it still exists; this method picks up the case where the
+    /// deleted entry **was** the selected one, by falling back to the same
+    /// numeric row clamped to the new `unified_results` length.
+    pub fn restore_selection_after_delete(&mut self, previous_index: Option<usize>) {
+        if self.selected_index.is_some() {
+            return;
+        }
+        self.selected_index = crate::session_list::clamp_index_after_delete(
+            previous_index,
+            self.unified_results.len(),
+        );
+    }
+
     fn collect_all_sessions(
         active_sessions: &[SessionUiInfo],
         resurrectable_sessions: &[(String, Duration)],