re-run command on enter and close pane on ctrl-c

10 files changed, 206 insertions, 63 deletions

diff --git a/zellij-server/src/os_input_output.rs b/zellij-server/src/os_input_output.rs
index 9bcc5244c..d1de70331 100644
--- a/zellij-server/src/os_input_output.rs
+++ b/zellij-server/src/os_input_output.rs
@@ -337,6 +337,13 @@ pub trait ServerOsApi: Send + Sync {
     fn get_cwd(&self, pid: Pid) -> Option<PathBuf>;
     /// Writes the given buffer to a string
     fn write_to_file(&mut self, buf: String, file: Option<String>);
+
+    fn re_run_command_in_terminal(
+        &self,
+        terminal_id: u32,
+        run_command: RunCommand,
+        quit_cb: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>, // u32 is the exit status
+    ) -> Result<(RawFd, RawFd), &'static str>;
 }
 
 impl ServerOsApi for ServerOsInputOutput {
@@ -485,6 +492,28 @@ impl ServerOsApi for ServerOsInputOutput {
             log::error!("could not write to file: {}", e);
         }
     }
+    fn re_run_command_in_terminal(
+        &self,
+        terminal_id: u32,
+        run_command: RunCommand,
+        quit_cb: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>, // u32 is the exit status
+    ) -> Result<(RawFd, RawFd), &'static str> {
+        let orig_termios = self.orig_termios.lock().unwrap();
+        let default_editor = None; // no need for a default editor when running an explicit command
+        match spawn_terminal(
+            TerminalAction::RunCommand(run_command),
+            orig_termios.clone(),
+            quit_cb,
+            default_editor,
+            terminal_id,
+        ) {
+            Ok((pid_primary, pid_secondary)) => {
+                self.terminal_id_to_raw_fd.lock().unwrap().insert(terminal_id, pid_primary);
+                Ok((pid_primary, pid_secondary))
+            },
+            Err(e) => Err(e)
+        }
+    }
 }
 
 impl Clone for Box<dyn ServerOsApi> {
diff --git a/zellij-server/src/panes/grid.rs b/zellij-server/src/panes/grid.rs
index 677eed163..12849f813 100644
--- a/zellij-server/src/panes/grid.rs
+++ b/zellij-server/src/panes/grid.rs
@@ -1504,7 +1504,7 @@ impl Grid {
     pub fn mark_for_rerender(&mut self) {
         self.should_render = true;
     }
-    fn reset_terminal_state(&mut self) {
+    pub fn reset_terminal_state(&mut self) {
         self.lines_above = VecDeque::with_capacity(*SCROLL_BUFFER_SIZE.get().unwrap());
         self.lines_below = vec![];
         self.viewport = vec![Row::new(self.width).canonical()];
diff --git a/zellij-server/src/panes/plugin_pane.rs b/zellij-server/src/panes/plugin_pane.rs
index 2ee865fe0..2bb781548 100644
--- a/zellij-server/src/panes/plugin_pane.rs
+++ b/zellij-server/src/panes/plugin_pane.rs
@@ -113,10 +113,6 @@ impl Pane for PluginPane {
     fn cursor_coordinates(&self) -> Option<(usize, usize)> {
         None
     }
-    fn adjust_input_to_terminal(&self, _input_bytes: Vec<u8>) -> Vec<u8> {
-        // noop
-        vec![]
-    }
     fn position_and_size(&self) -> PaneGeom {
         self.geom
     }
diff --git a/zellij-server/src/panes/terminal_pane.rs b/zellij-server/src/panes/terminal_pane.rs
index 9db5ff9d2..03dd0c878 100644
--- a/zellij-server/src/panes/terminal_pane.rs
+++ b/zellij-server/src/panes/terminal_pane.rs
@@ -6,7 +6,7 @@ use crate::panes::{
 };
 use crate::panes::{AnsiCode, LinkHandler};
 use crate::pty::VteBytes;
-use crate::tab::Pane;
+use crate::tab::{Pane, AdjustedInput};
 use crate::ClientId;
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
@@ -37,6 +37,8 @@ const HOME_KEY: &[u8] = &[27, 91, 72];
 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: &[u8] = &[13]; // TODO: check this to be sure it fits all types of ENTER
+const CTRL_C: &[u8] = &[3]; // TODO: check this to be sure it fits all types of CTRL_C (with mac, etc)
 const TERMINATING_STRING: &str = "\0";
 const DELETE_KEY: &str = "\u{007F}";
 const BACKSPACE_KEY: &str = "\u{0008}";
@@ -159,53 +161,69 @@ impl Pane for TerminalPane {
             .cursor_coordinates()
             .map(|(x, y)| (x + left, y + top))
     }
-    fn adjust_input_to_terminal(&self, input_bytes: Vec<u8>) -> Vec<u8> {
+    fn adjust_input_to_terminal(&mut self, input_bytes: Vec<u8>) -> Option<AdjustedInput> {
         // there are some cases in which the terminal state means that input sent to it
         // needs to be adjusted.
         // here we match against those cases - if need be, we adjust the input and if not
         // we send back the original input
-        if self.grid.new_line_mode {
-            if let &[13] = input_bytes.as_slice() {
-                // LNM - carriage return is followed by linefeed
-                return "\u{0d}\u{0a}".as_bytes().to_vec();
-            };
-        }
-        if self.grid.cursor_key_mode {
+        if let Some((_exit_status, run_command)) = &self.is_held {
             match input_bytes.as_slice() {
-                LEFT_ARROW => {
-                    return AnsiEncoding::Left.as_vec_bytes();
-                },
-                RIGHT_ARROW => {
-                    return AnsiEncoding::Right.as_vec_bytes();
-                },
-                UP_ARROW => {
-                    return AnsiEncoding::Up.as_vec_bytes();
-                },
-                DOWN_ARROW => {
-                    return AnsiEncoding::Down.as_vec_bytes();
-                },
-
-                HOME_KEY => {
-                    return AnsiEncoding::Home.as_vec_bytes();
-                },
-                END_KEY => {
-                    return AnsiEncoding::End.as_vec_bytes();
-                },
-                _ => {},
-            };
-        }
+                ENTER => {
+                    let run_command = run_command.clone();
+                    self.is_held = None;
+                    self.grid.reset_terminal_state();
+                    self.set_should_render(true);
+                    Some(AdjustedInput::ReRunCommandInThisPane(run_command))
+                }
+                CTRL_C => {
+                    Some(AdjustedInput::CloseThisPane)
+                }
+                _ => None
+            }
+        } else {
+            if self.grid.new_line_mode {
+                if let &[13] = input_bytes.as_slice() {
+                    // LNM - carriage return is followed by linefeed
+                    return Some(AdjustedInput::WriteBytesToTerminal("\u{0d}\u{0a}".as_bytes().to_vec()));
+                };
+            }
+            if self.grid.cursor_key_mode {
+                match input_bytes.as_slice() {
+                    LEFT_ARROW => {
+                        return Some(AdjustedInput::WriteBytesToTerminal(AnsiEncoding::Left.as_vec_bytes()));
+                    },
+                    RIGHT_ARROW => {
+                        return Some(AdjustedInput::WriteBytesToTerminal(AnsiEncoding::Right.as_vec_bytes()));
+                    },
+                    UP_ARROW => {
+                        return Some(AdjustedInput::WriteBytesToTerminal(AnsiEncoding::Up.as_vec_bytes()));
+                    },
+                    DOWN_ARROW => {
+                        return Some(AdjustedInput::WriteBytesToTerminal(AnsiEncoding::Down.as_vec_bytes()));
+                    },
+
+                    HOME_KEY => {
+                        return Some(AdjustedInput::WriteBytesToTerminal(AnsiEncoding::Home.as_vec_bytes()));
+                    },
+                    END_KEY => {
+                        return Some(AdjustedInput::WriteBytesToTerminal(AnsiEncoding::End.as_vec_bytes()));
+                    },
+                    _ => {},
+                };
+            }
 
-        if !self.grid.bracketed_paste_mode {
-            // Zellij itself operates in bracketed paste mode, so the terminal sends these
-            // instructions (bracketed paste start and bracketed paste end respectively)
-            // when pasting input. We only need to make sure not to send them to terminal
-            // panes who do not work in this mode
-            match input_bytes.as_slice() {
-                BRACKETED_PASTE_BEGIN | BRACKETED_PASTE_END => return vec![],
-                _ => {},
+            if !self.grid.bracketed_paste_mode {
+                // Zellij itself operates in bracketed paste mode, so the terminal sends these
+                // instructions (bracketed paste start and bracketed paste end respectively)
+                // when pasting input. We only need to make sure not to send them to terminal
+                // panes who do not work in this mode
+                match input_bytes.as_slice() {
+                    BRACKETED_PASTE_BEGIN | BRACKETED_PASTE_END => return Some(AdjustedInput::WriteBytesToTerminal(vec![])),
+                    _ => {},
+                }
             }
+            Some(AdjustedInput::WriteBytesToTerminal(input_bytes))
         }
-        input_bytes
     }
     fn position_and_size(&self) -> PaneGeom {
         self.geom
diff --git a/zellij-server/src/pty.rs b/zellij-server/src/pty.rs
index 7114b9513..7c07c7866 100644
--- a/zellij-server/src/pty.rs
+++ b/zellij-server/src/pty.rs
@@ -41,6 +41,7 @@ pub(crate) enum PtyInstruction {
     ), // the String is the tab name
     ClosePane(PaneId),
     CloseTab(Vec<PaneId>),
+    ReRunCommandInPane(PaneId, RunCommand),
     Exit,
 }
 
@@ -56,6 +57,7 @@ impl From<&PtyInstruction> for PtyContext {
             PtyInstruction::ClosePane(_) => PtyContext::ClosePane,
             PtyInstruction::CloseTab(_) => PtyContext::CloseTab,
             PtyInstruction::NewTab(..) => PtyContext::NewTab,
+            PtyInstruction::ReRunCommandInPane(..) => PtyContext::ReRunCommandInPane,
             PtyInstruction::Exit => PtyContext::Exit,
         }
     }
@@ -175,6 +177,11 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) {
                     .send_to_server(ServerInstruction::UnblockInputThread)
                     .unwrap();
             },
+            PtyInstruction::ReRunCommandInPane(pane_id, run_command) => {
+                if let Err(e) = pty.rerun_command_in_pane(pane_id, run_command) {
+                    log::error!("Failed to rerun command in pane: {:?}", e);
+                }
+            }
             PtyInstruction::Exit => break,
         }
     }
@@ -382,6 +389,45 @@ impl Pty {
             self.active_panes.insert(client_id, pane_id);
         }
     }
+    pub fn rerun_command_in_pane(&mut self, pane_id: PaneId, run_command: RunCommand) -> Result<(), &'static str> {
+        match pane_id {
+            PaneId::Terminal(id) => {
+                let _ = self.task_handles.remove(&id); // if all is well, this shouldn't be here
+                let _ = self.id_to_child_pid.remove(&id); // if all is wlel, this shouldn't be here
+
+                let quit_cb = Box::new({
+                    let senders = self.bus.senders.clone();
+                    move |pane_id, exit_status, command| {
+                        // we only re-run held panes, so we'll never close them from Pty
+                        let _ = senders.send_to_screen(ScreenInstruction::HoldPane(pane_id, exit_status, command, None));
+                    }
+                });
+                let (pid_primary, child_fd): (RawFd, RawFd) = self
+                    .bus
+                    .os_input
+                    .as_mut()
+                    .unwrap()
+                    .re_run_command_in_terminal(id, run_command, quit_cb)?;
+                let terminal_bytes = task::spawn({
+                    let senders = self.bus.senders.clone();
+                    let os_input = self.bus.os_input.as_ref().unwrap().clone();
+                    let debug_to_file = self.debug_to_file;
+                    async move {
+                        TerminalBytes::new(pid_primary, senders, os_input, debug_to_file, id)
+                            .listen()
+                            .await;
+                    }
+                });
+
+                self.task_handles.insert(id, terminal_bytes);
+                self.id_to_child_pid.insert(id, child_fd);
+                Ok(())
+            },
+            _ => {
+                Err("Tried to respawn a plugin pane - cannot respawn plugin panes")
+            }
+        }
+    }
 }
 
 impl Drop for Pty {
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index 819c14e92..a0b216f01 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -99,6 +99,7 @@ pub(crate) struct Tab {
     last_mouse_hold_position: Option<Position>,
     terminal_emulator_colors: Rc<RefCell<Palette>>,
     terminal_emulator_color_codes: Rc<RefCell<HashMap<usize, String>>>,
+    pids_waiting_resize: HashSet<u32>, // u32 is the terminal_id
 }
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
@@ -126,7 +127,9 @@ pub trait Pane {
     fn set_geom_override(&mut self, pane_geom: PaneGeom);
     fn handle_pty_bytes(&mut self, bytes: VteBytes);
     fn cursor_coordinates(&self) -> Option<(usize, usize)>;
-    fn adjust_input_to_terminal(&self, input_bytes: Vec<u8>) -> Vec<u8>;
+    fn adjust_input_to_terminal(&mut self, input_bytes: Vec<u8>) -> Option<AdjustedInput> {
+        None
+    }
     fn position_and_size(&self) -> PaneGeom;
     fn current_geom(&self) -> PaneGeom;
     fn geom_override(&self) -> Option<PaneGeom>;
@@ -355,6 +358,13 @@ pub trait Pane {
     }
 }
 
+#[derive(Clone, Debug)]
+pub enum AdjustedInput {
+    WriteBytesToTerminal(Vec<u8>),
+    ReRunCommandInThisPane(RunCommand),
+    CloseThisPane,
+}
+
 impl Tab {
     // FIXME: Still too many arguments for clippy to be happy...
     #[allow(clippy::too_many_arguments)]
@@ -452,6 +462,7 @@ impl Tab {
             last_mouse_hold_position: None,
             terminal_emulator_colors,
             terminal_emulator_color_codes,
+            pids_waiting_resize: HashSet::new(),
         }
     }
 
@@ -995,6 +1006,9 @@ impl Tab {
                     .find(|s_p| s_p.pid() == PaneId::Terminal(pid))
             })
         {
+            if self.pids_waiting_resize.remove(&pid) {
+                resize_pty!(terminal_output, self.os_api);
+            }
             terminal_output.handle_pty_bytes(bytes);
             let messages_to_pty = terminal_output.drain_messages_to_pty();
             let clipboard_update = terminal_output.drain_clipboard_update();
@@ -1042,18 +1056,33 @@ impl Tab {
             PaneId::Terminal(active_terminal_id) => {
                 let active_terminal = self
                     .floating_panes
-                    .get(&pane_id)
-                    .or_else(|| self.tiled_panes.get_pane(pane_id))
-                    .or_else(|| self.suppressed_panes.get(&pane_id))
-                    .unwrap();
-                let adjusted_input = active_terminal.adjust_input_to_terminal(input_bytes);
-
-                self.senders
-                    .send_to_pty_writer(PtyWriteInstruction::Write(
-                        adjusted_input,
-                        active_terminal_id,
-                    ))
+                    .get_mut(&pane_id)
+                    .or_else(|| self.tiled_panes.get_pane_mut(pane_id))
+                    .or_else(|| self.suppressed_panes.get_mut(&pane_id))
                     .unwrap();
+                match active_terminal.adjust_input_to_terminal(input_bytes) {
+                    Some(AdjustedInput::WriteBytesToTerminal(adjusted_input)) => {
+                        self.senders
+                            .send_to_pty_writer(PtyWriteInstruction::Write(
+                                adjusted_input,
+                                active_terminal_id,
+                            ))
+                            .unwrap();
+                    },
+                    Some(AdjustedInput::ReRunCommandInThisPane(command)) => {
+                        self.pids_waiting_resize.insert(active_terminal_id);
+                        self.senders
+                            .send_to_pty(PtyInstruction::ReRunCommandInPane(
+                                PaneId::Terminal(active_terminal_id),
+                                command,
+                            ))
+                            .unwrap();
+                    },
+                    Some(AdjustedInput::CloseThisPane) => {
+                        self.close_pane(PaneId::Terminal(active_terminal_id), false);
+                    },
+                    None => {}
+                }
             },
             PaneId::Plugin(pid) => {
                 for key in parse_keys(&input_bytes) {
diff --git a/zellij-server/src/tab/unit/tab_integration_tests.rs b/zellij-server/src/tab/unit/tab_integration_tests.rs
index 1eb6f2699..026646e69 100644
--- a/zellij-server/src/tab/unit/tab_integration_tests.rs
+++ b/zellij-server/src/tab/unit/tab_integration_tests.rs
@@ -30,7 +30,7 @@ use zellij_utils::nix;
 
 use zellij_utils::{
     data::{InputMode, ModeInfo, Palette, Style},
-    input::command::TerminalAction,
+    input::command::{TerminalAction, RunCommand},
     interprocess::local_socket::LocalSocketStream,
     ipc::{ClientToServerMsg, ServerToClientMsg},
 };
@@ -47,7 +47,7 @@ impl ServerOsApi for FakeInputOutput {
     fn spawn_terminal(
         &self,
         _file_to_open: TerminalAction,
-        _quit_cb: Box<dyn Fn(PaneId) + Send>,
+        _quit_db: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>,
         _default_editor: Option<PathBuf>,
     ) -> Result<(u32, RawFd, RawFd), &'static str> {
         unimplemented!()
@@ -103,6 +103,14 @@ impl ServerOsApi for FakeInputOutput {
         };
         self.file_dumps.lock().unwrap().insert(f, buf);
     }
+    fn re_run_command_in_terminal(
+        &self,
+        _terminal_id: u32,
+        _run_command: RunCommand,
+        _quit_cb: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>, // u32 is the exit status
+    ) -> Result<(RawFd, RawFd), &'static str> {
+        unimplemented!()
+    }
 }
 
 struct MockPtyInstructionBus {
diff --git a/zellij-server/src/tab/unit/tab_tests.rs b/zellij-server/src/tab/unit/tab_tests.rs
index 23d25c058..635b710ab 100644
--- a/zellij-server/src/tab/unit/tab_tests.rs
+++ b/zellij-server/src/tab/unit/tab_tests.rs
@@ -21,7 +21,7 @@ use zellij_utils::nix;
 
 use zellij_utils::{
     data::{ModeInfo, Palette, Style},
-    input::command::TerminalAction,
+    input::command::{TerminalAction, RunCommand},
     interprocess::local_socket::LocalSocketStream,
     ipc::{ClientToServerMsg, ServerToClientMsg},
 };
@@ -36,7 +36,7 @@ impl ServerOsApi for FakeInputOutput {
     fn spawn_terminal(
         &self,
         _file_to_open: TerminalAction,
-        _quit_cb: Box<dyn Fn(PaneId) + Send>,
+        _quit_cb: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>,
         _default_editor: Option<PathBuf>,
     ) -> Result<(u32, RawFd, RawFd), &'static str> {
         unimplemented!()
@@ -89,6 +89,14 @@ impl ServerOsApi for FakeInputOutput {
     fn write_to_file(&mut self, _buf: String, _name: Option<String>) {
         unimplemented!()
     }
+    fn re_run_command_in_terminal(
+        &self,
+        _terminal_id: u32,
+        _run_command: RunCommand,
+        _quit_cb: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>, // u32 is the exit status
+    ) -> Result<(RawFd, RawFd), &'static str> {
+        unimplemented!()
+    }
 }
 
 fn create_new_tab(size: Size) -> Tab {
diff --git a/zellij-server/src/unit/screen_tests.rs b/zellij-server/src/unit/screen_tests.rs
index 036d29c52..9b953bd68 100644
--- a/zellij-server/src/unit/screen_tests.rs
+++ b/zellij-server/src/unit/screen_tests.rs
@@ -12,7 +12,7 @@ use std::path::PathBuf;
 use zellij_utils::cli::CliAction;
 use zellij_utils::errors::ErrorContext;
 use zellij_utils::input::actions::{Action, Direction, ResizeDirection};
-use zellij_utils::input::command::TerminalAction;
+use zellij_utils::input::command::{TerminalAction, RunCommand};
 use zellij_utils::input::layout::{PaneLayout, SplitDirection};
 use zellij_utils::input::options::Options;
 use zellij_utils::ipc::IpcReceiverWithContext;
@@ -132,7 +132,7 @@ impl ServerOsApi for FakeInputOutput {
     fn spawn_terminal(
         &self,
         _file_to_open: TerminalAction,
-        _quit_db: Box<dyn Fn(PaneId) + Send>,
+        _quit_db: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>,
         _default_editor: Option<PathBuf>,
     ) -> Result<(u32, RawFd, RawFd), &'static str> {
         unimplemented!()
@@ -195,6 +195,14 @@ impl ServerOsApi for FakeInputOutput {
                 .insert(filename, contents);
         }
     }
+    fn re_run_command_in_terminal(
+        &self,
+        _terminal_id: u32,
+        _run_command: RunCommand,
+        _quit_cb: Box<dyn Fn(PaneId, Option<i32>, RunCommand) + Send>, // u32 is the exit status
+    ) -> Result<(RawFd, RawFd), &'static str> {
+        unimplemented!()
+    }
 }
 
 fn create_new_screen(size: Size) -> Screen {
diff --git a/zellij-utils/src/errors.rs b/zellij-utils/src/errors.rs
index 96663bc27..30b2e1300 100644
--- a/zellij-utils/src/errors.rs
+++ b/zellij-utils/src/errors.rs
@@ -317,6 +317,7 @@ pub enum PtyContext {
     NewTab,
     ClosePane,
     CloseTab,
+    ReRunCommandInPane,
     Exit,
 }