feat(panes): Add an option to press <ESC> and drop to shell in command panes (#2872)

* feat(panes): ESC to drop to default shell on command panes

* style(fmt): rustfmt

16 files changed, 161 insertions, 50 deletions

diff --git a/zellij-server/src/os_input_output.rs b/zellij-server/src/os_input_output.rs
index 46a0d9974..793e5b665 100644
--- a/zellij-server/src/os_input_output.rs
+++ b/zellij-server/src/os_input_output.rs
@@ -31,7 +31,7 @@ use zellij_utils::{
 };
 
 use std::{
-    collections::{BTreeMap, HashMap, HashSet},
+    collections::{BTreeMap, BTreeSet, HashMap},
     env,
     fs::File,
     io::Write,
@@ -581,7 +581,7 @@ impl ServerOsApi for ServerOsInputOutput {
             .with_context(err_context)?;
         let mut terminal_id = None;
         {
-            let current_ids: HashSet<u32> = self
+            let current_ids: BTreeSet<u32> = self
                 .terminal_id_to_raw_fd
                 .lock()
                 .to_anyhow()
@@ -589,13 +589,7 @@ impl ServerOsApi for ServerOsInputOutput {
                 .keys()
                 .copied()
                 .collect();
-            for i in 0..u32::MAX {
-                let i = i as u32;
-                if !current_ids.contains(&i) {
-                    terminal_id = Some(i);
-                    break;
-                }
-            }
+            terminal_id = current_ids.last().map(|l| l + 1).or(Some(0));
         }
         match terminal_id {
             Some(terminal_id) => {
@@ -628,7 +622,7 @@ impl ServerOsApi for ServerOsInputOutput {
 
         let mut terminal_id = None;
         {
-            let current_ids: HashSet<u32> = self
+            let current_ids: BTreeSet<u32> = self
                 .terminal_id_to_raw_fd
                 .lock()
                 .to_anyhow()
@@ -636,13 +630,7 @@ impl ServerOsApi for ServerOsInputOutput {
                 .keys()
                 .copied()
                 .collect();
-            for i in 0..u32::MAX {
-                let i = i as u32;
-                if !current_ids.contains(&i) {
-                    terminal_id = Some(i);
-                    break;
-                }
-            }
+            terminal_id = current_ids.last().map(|l| l + 1).or(Some(0));
         }
         match terminal_id {
             Some(terminal_id) => {
diff --git a/zellij-server/src/panes/terminal_character.rs b/zellij-server/src/panes/terminal_character.rs
index 23ea1a524..71aec7a45 100644
--- a/zellij-server/src/panes/terminal_character.rs
+++ b/zellij-server/src/panes/terminal_character.rs
@@ -773,21 +773,25 @@ pub fn render_first_run_banner(
 
             let controls_bare_text_first_part = "<";
             let enter_bare_text = "ENTER";
-            let controls_bare_text_second_part = "> to run, <";
+            let controls_bare_text_second_part = "> run, <";
+            let esc_bare_text = "ESC";
+            let controls_bare_text_third_part = "> drop to shell, <";
             let ctrl_c_bare_text = "Ctrl-c";
-            let controls_bare_text_third_part = "> to exit";
+            let controls_bare_text_fourth_part = "> exit";
             let controls_color = RESET_STYLES
                 .foreground(Some(AnsiCode::from(style.colors.orange)))
                 .bold(Some(AnsiCode::On));
             let controls_line_length = controls_bare_text_first_part.len()
                 + enter_bare_text.len()
                 + controls_bare_text_second_part.len()
+                + esc_bare_text.len()
+                + controls_bare_text_third_part.len()
                 + ctrl_c_bare_text.len()
-                + controls_bare_text_third_part.len();
+                + controls_bare_text_fourth_part.len();
             let controls_column_start_position =
                 middle_column.saturating_sub(controls_line_length / 2);
             let controls_line = format!(
-                "\u{1b}[{};{}H{}<{}{}{}{}> to run, <{}{}{}{}> to exit",
+                "\u{1b}[{};{}H{}<{}{}{}{}> run, <{}{}{}{}> drop to shell, <{}{}{}{}> exit",
                 middle_row + 2,
                 controls_column_start_position,
                 bold_text,
@@ -796,6 +800,10 @@ pub fn render_first_run_banner(
                 RESET_STYLES,
                 bold_text,
                 controls_color,
+                esc_bare_text,
+                RESET_STYLES,
+                bold_text,
+                controls_color,
                 ctrl_c_bare_text,
                 RESET_STYLES,
                 bold_text
@@ -817,21 +825,25 @@ pub fn render_first_run_banner(
 
             let controls_bare_text_first_part = "<";
             let enter_bare_text = "ENTER";
-            let controls_bare_text_second_part = "> to run, <";
+            let controls_bare_text_second_part = "> run, <";
+            let esc_bare_text = "ESC";
+            let controls_bare_text_third_part = "> drop to shell, <";
             let ctrl_c_bare_text = "Ctrl-c";
-            let controls_bare_text_third_part = "> to exit";
+            let controls_bare_text_fourth_part = "> exit";
             let controls_color = RESET_STYLES
                 .foreground(Some(AnsiCode::from(style.colors.orange)))
                 .bold(Some(AnsiCode::On));
             let controls_line_length = controls_bare_text_first_part.len()
                 + enter_bare_text.len()
                 + controls_bare_text_second_part.len()
+                + esc_bare_text.len()
+                + controls_bare_text_third_part.len()
                 + ctrl_c_bare_text.len()
-                + controls_bare_text_third_part.len();
+                + controls_bare_text_fourth_part.len();
             let controls_column_start_position =
                 middle_column.saturating_sub(controls_line_length / 2);
             let controls_line = format!(
-                "\u{1b}[{};{}H{}<{}{}{}{}> to run, <{}{}{}{}> to exit",
+                "\u{1b}[{};{}H{}<{}{}{}{}> run, <{}{}{}{}> drop to shell, <{}{}{}{}> exit",
                 middle_row + 2,
                 controls_column_start_position,
                 bold_text,
@@ -840,6 +852,10 @@ pub fn render_first_run_banner(
                 RESET_STYLES,
                 bold_text,
                 controls_color,
+                esc_bare_text,
+                RESET_STYLES,
+                bold_text,
+                controls_color,
                 ctrl_c_bare_text,
                 RESET_STYLES,
                 bold_text
diff --git a/zellij-server/src/panes/terminal_pane.rs b/zellij-server/src/panes/terminal_pane.rs
index 6c405b6e5..ac4c6da63 100644
--- a/zellij-server/src/panes/terminal_pane.rs
+++ b/zellij-server/src/panes/terminal_pane.rs
@@ -40,6 +40,7 @@ const END_KEY: &[u8] = &[27, 91, 70];
 const BRACKETED_PASTE_BEGIN: &[u8] = &[27, 91, 50, 48, 48, 126];
 const BRACKETED_PASTE_END: &[u8] = &[27, 91, 50, 48, 49, 126];
 const ENTER_NEWLINE: &[u8] = &[10];
+const ESC: &[u8] = &[27];
 const ENTER_CARRIAGE_RETURN: &[u8] = &[13];
 const SPACE: &[u8] = &[32];
 const CTRL_C: &[u8] = &[3]; // TODO: check this to be sure it fits all types of CTRL_C (with mac, etc)
@@ -192,6 +193,13 @@ impl Pane for TerminalPane {
                     self.remove_banner();
                     Some(AdjustedInput::ReRunCommandInThisPane(run_command))
                 },
+                ESC => {
+                    self.is_held = None;
+                    self.grid.reset_terminal_state();
+                    self.set_should_render(true);
+                    self.remove_banner();
+                    Some(AdjustedInput::DropToShellInThisPane)
+                },
                 CTRL_C => Some(AdjustedInput::CloseThisPane),
                 _ => None,
             }
diff --git a/zellij-server/src/pty.rs b/zellij-server/src/pty.rs
index e355e8c55..6dd38e9fe 100644
--- a/zellij-server/src/pty.rs
+++ b/zellij-server/src/pty.rs
@@ -66,6 +66,7 @@ pub enum PtyInstruction {
     ClosePane(PaneId),
     CloseTab(Vec<PaneId>),
     ReRunCommandInPane(PaneId, RunCommand),
+    DropToShellInPane(PaneId, Option<PathBuf>), // Option<PathBuf> - default shell
     SpawnInPlaceTerminal(
         Option<TerminalAction>,
         Option<String>,
@@ -89,6 +90,7 @@ impl From<&PtyInstruction> for PtyContext {
             PtyInstruction::CloseTab(_) => PtyContext::CloseTab,
             PtyInstruction::NewTab(..) => PtyContext::NewTab,
             PtyInstruction::ReRunCommandInPane(..) => PtyContext::ReRunCommandInPane,
+            PtyInstruction::DropToShellInPane(..) => PtyContext::DropToShellInPane,
             PtyInstruction::SpawnInPlaceTerminal(..) => PtyContext::SpawnInPlaceTerminal,
             PtyInstruction::DumpLayout(..) => PtyContext::DumpLayout,
             PtyInstruction::LogLayoutToHd(..) => PtyContext::LogLayoutToHd,
@@ -532,6 +534,55 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) -> Result<()> {
                     },
                 }
             },
+            PtyInstruction::DropToShellInPane(pane_id, default_shell) => {
+                let err_context = || format!("failed to rerun command in pane {:?}", pane_id);
+
+                // TODO: get configured default_shell from screen/tab as an option and default to
+                // this otherwise (also look for a place that turns get_default_shell into a
+                // RunCommand, we might have done this before)
+                let run_command = RunCommand {
+                    command: default_shell.unwrap_or_else(|| get_default_shell()),
+                    hold_on_close: false,
+                    hold_on_start: false,
+                    // TODO: cwd
+                    ..Default::default()
+                };
+                match pty
+                    .rerun_command_in_pane(pane_id, run_command.clone())
+                    .with_context(err_context)
+                {
+                    Ok(..) => {},
+                    Err(err) => match err.downcast_ref::<ZellijError>() {
+                        Some(ZellijError::CommandNotFound { terminal_id, .. }) => {
+                            if run_command.hold_on_close {
+                                pty.bus
+                                    .senders
+                                    .send_to_screen(ScreenInstruction::PtyBytes(
+                                        *terminal_id,
+                                        format!(
+                                            "Command not found: {}",
+                                            run_command.command.display()
+                                        )
+                                        .as_bytes()
+                                        .to_vec(),
+                                    ))
+                                    .with_context(err_context)?;
+                                pty.bus
+                                    .senders
+                                    .send_to_screen(ScreenInstruction::HoldPane(
+                                        PaneId::Terminal(*terminal_id),
+                                        Some(2), // exit status
+                                        run_command,
+                                        None,
+                                        None,
+                                    ))
+                                    .with_context(err_context)?;
+                            }
+                        },
+                        _ => Err::<(), _>(err).non_fatal(),
+                    },
+                }
+            },
             PtyInstruction::DumpLayout(mut session_layout_metadata, client_id) => {
                 let err_context = || format!("Failed to dump layout");
                 pty.populate_session_layout_metadata(&mut session_layout_metadata);
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index 6b974ef85..582861856 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -1149,6 +1149,7 @@ impl Screen {
             self.terminal_emulator_colors.clone(),
             self.terminal_emulator_color_codes.clone(),
             swap_layouts,
+            self.default_shell.clone(),
             self.debug,
         );
         self.tabs.insert(tab_index, tab);
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index 30a6bed8c..f441e7292 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -8,6 +8,7 @@ mod swap_layouts;
 
 use copy_command::CopyCommand;
 use std::env::temp_dir;
+use std::path::PathBuf;
 use uuid::Uuid;
 use zellij_utils::data::{
     Direction, PaneInfo, PermissionStatus, PermissionType, PluginPermission, ResizeStrategy,
@@ -182,6 +183,7 @@ pub(crate) struct Tab {
     pending_instructions: Vec<BufferedTabInstruction>, // instructions that came while the tab was
     // pending and need to be re-applied
     swap_layouts: SwapLayouts,
+    default_shell: Option<PathBuf>,
     debug: bool,
 }
 
@@ -482,6 +484,7 @@ pub enum AdjustedInput {
     ReRunCommandInThisPane(RunCommand),
     PermissionRequestResult(Vec<PermissionType>, PermissionStatus),
     CloseThisPane,
+    DropToShellInThisPane,
 }
 pub fn get_next_terminal_position(
     tiled_panes: &TiledPanes,
@@ -528,6 +531,7 @@ impl Tab {
         terminal_emulator_colors: Rc<RefCell<Palette>>,
         terminal_emulator_color_codes: Rc<RefCell<HashMap<usize, String>>>,
         swap_layouts: (Vec<SwapTiledLayout>, Vec<SwapFloatingLayout>),
+        default_shell: Option<PathBuf>,
         debug: bool,
     ) -> Self {
         let name = if name.is_empty() {
@@ -615,6 +619,7 @@ impl Tab {
             is_pending: true, // will be switched to false once the layout is applied
             pending_instructions: vec![],
             swap_layouts,
+            default_shell,
             debug,
         }
     }
@@ -1695,6 +1700,16 @@ impl Tab {
                         self.close_pane(PaneId::Terminal(active_terminal_id), false, None);
                         should_update_ui = true;
                     },
+                    Some(AdjustedInput::DropToShellInThisPane) => {
+                        self.pids_waiting_resize.insert(active_terminal_id);
+                        self.senders
+                            .send_to_pty(PtyInstruction::DropToShellInPane(
+                                PaneId::Terminal(active_terminal_id),
+                                self.default_shell.clone(),
+                            ))
+                            .with_context(err_context)?;
+                        should_update_ui = true;
+                    },
                     Some(_) => {},
                     None => {},
                 }
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_floating_layout_is_included_in_swap_layouts.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_floating_layout_is_included_in_swap_layouts.snap
index f6e332910..cd3f576f9 100644
--- a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_floating_layout_is_included_in_swap_layouts.snap
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_floating_layout_is_included_in_swap_layouts.snap
@@ -1,6 +1,6 @@
 ---
 source: zellij-server/src/tab/./unit/tab_integration_tests.rs
-assertion_line: 5472
+assertion_line: 6071
 expression: snapshot
 ---
 00 (C): ┌ tab-bar ────────────────────────────────────────┌ status-bar ──────────────────────────────────────────────┐──────────┐
@@ -19,7 +19,7 @@ expression: snapshot
 13 (C): │                             │                   │                                                          │          │
 14 (C): │                Waiting to ru└───────────────────│                Waiting to run: command2                  │          │
 15 (C): │                                                 │                                                          │          │
-16 (C): │            <ENTER> to run, <Ctrl-c> to exit     │            <ENTER> to run, <Ctrl-c> to exit              │          │
+16 (C): │     <ENTER> run, <ESC> drop to shell, <Ctrl-c> e│     <ENTER> run, <ESC> drop to shell, <Ctrl-c> exit      │          │
 17 (C): │                                                 │                                                          │          │
 18 (C): │                                                 │                                                          │          │
 19 (C): └─────────────────────────────────────────────────└──────────────────────────────────────────────────────────┘──────────┘
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_layout_is_included_in_swap_layouts.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_layout_is_included_in_swap_layouts.snap
index 97692e660..440d15cb8 100644
--- a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_layout_is_included_in_swap_layouts.snap
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__base_layout_is_included_in_swap_layouts.snap
@@ -1,5 +1,6 @@
 ---
 source: zellij-server/src/tab/./unit/tab_integration_tests.rs
+assertion_line: 5109
 expression: snapshot
 ---
 00 (C): I am a tab bar                                                                                                           
@@ -12,14 +13,14 @@ expression: snapshot
 07 (C): │                                       ││                                      ││                                      │
 08 (C): │      Waiting to run: command2         ││                                      ││      Waiting to run: command1        │
 09 (C): │                                       ││                                      ││                                      │
-10 (C): │  <ENTER> to run, <Ctrl-c> to exit     ││                                      ││  <ENTER> to run, <Ctrl-c> to exit    │
-11 (C): │                                       ││                                      ││                                      │
+10 (C): │<ENTER> run, <ESC> drop to shell, <Ctrl││                                      ││<ENTER> run, <ESC> drop to shell, <Ctr│
+11 (C): │-c> exit                               ││                                      ││l-c> exit                             │
 12 (C): │                                       ││                                      ││                                      │
 13 (C): │                                       ││                                      ││                                      │
 14 (C): │                                       ││                                      ││                                      │
 15 (C): │                                       ││                                      ││                                      │
 16 (C): │                                       ││                                      ││                                      │
-17 (C): └───────────────────────────────────────┘└──────────────────────────────────────┘└─ <ENTER> to run, <Ctrl-c> to exit ───┘
+17 (C): └───────────────────────────────────────┘└──────────────────────────────────────┘└──────────────────────────────────────┘
 18 (C): I am a                                                                                                                   
 19 (C): status bar                                                                                                               
 
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__floating_layout_with_plugins_and_commands_swaped_properly.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__floating_layout_with_plugins_and_commands_swaped_properly.snap
index 398383482..109824752 100644
--- a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__floating_layout_with_plugins_and_commands_swaped_properly.snap
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__floating_layout_with_plugins_and_commands_swaped_properly.snap
@@ -1,6 +1,6 @@
 ---
 source: zellij-server/src/tab/./unit/tab_integration_tests.rs
-assertion_line: 5926
+assertion_line: 5971
 expression: snapshot
 ---
 00 (C): ┌ status-bar ──────────────────────────────────────────────┐─────────────────────────────────────────────────┐──────────┐
@@ -19,7 +19,7 @@ expression: snapshot
 13 (C): │                             │                   │                                                          │          │
 14 (C): │                Waiting to ru└───────────────────│                Waiting to run: command2                  │          │
 15 (C): │                                                 │                                                          │          │
-16 (C): │            <ENTER> to run, <Ctrl-c> to exit     │            <ENTER> to run, <Ctrl-c> to exit              │          │
+16 (C): │     <ENTER> run, <ESC> drop to shell, <Ctrl-c> e│     <ENTER> run, <ESC> drop to shell, <Ctrl-c> exit      │          │
 17 (C): │                                                 │                                                          │          │
 18 (C): │                                                 │                                                          │