fix(messaging): cache hold pane messages by their tab_id if the tab is not ready (#2196)

* fix(messaging): cache hold pane messages by their tab_id if the tab is not ready

* style(fmt): rustfmt

3 files changed, 50 insertions, 7 deletions

diff --git a/zellij-server/src/pty.rs b/zellij-server/src/pty.rs
index fd67234c3..4c4f0a910 100644
--- a/zellij-server/src/pty.rs
+++ b/zellij-server/src/pty.rs
@@ -147,6 +147,7 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) -> Result<()> {
                                         pty.bus.senders.clone(),
                                         *terminal_id,
                                         run_command.clone(),
+                                        None,
                                     )
                                     .with_context(err_context)?;
                                 }
@@ -243,6 +244,7 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) -> Result<()> {
                                             Some(2), // exit status
                                             run_command,
                                             None,
+                                            None,
                                         ))
                                         .with_context(err_context)?;
                                 }
@@ -313,6 +315,7 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) -> Result<()> {
                                             Some(2), // exit status
                                             run_command,
                                             None,
+                                            None,
                                         ))
                                         .with_context(err_context)?;
                                 }
@@ -422,6 +425,7 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) -> Result<()> {
                                         Some(2), // exit status
                                         run_command,
                                         None,
+                                        None,
                                     ))
                                     .with_context(err_context)?;
                             }
@@ -532,6 +536,7 @@ impl Pty {
                         exit_status,
                         command,
                         None,
+                        None,
                     ));
                 } else {
                     let _ = senders.send_to_screen(ScreenInstruction::ClosePane(pane_id, None));
@@ -599,14 +604,14 @@ impl Pty {
                     // new_pane_pids
         for run_instruction in extracted_run_instructions {
             if let Some(new_pane_data) =
-                self.apply_run_instruction(run_instruction, default_shell.clone())?
+                self.apply_run_instruction(run_instruction, default_shell.clone(), tab_index)?
             {
                 new_pane_pids.push(new_pane_data);
             }
         }
         for run_instruction in extracted_floating_run_instructions {
             if let Some(new_pane_data) =
-                self.apply_run_instruction(run_instruction, default_shell.clone())?
+                self.apply_run_instruction(run_instruction, default_shell.clone(), tab_index)?
             {
                 new_floating_panes_pids.push(new_pane_data);
             }
@@ -686,6 +691,7 @@ impl Pty {
                                 self.bus.senders.clone(),
                                 terminal_id,
                                 run_command.clone(),
+                                Some(tab_index),
                             )
                             .with_context(err_context)?;
                         } else {
@@ -706,6 +712,7 @@ impl Pty {
         &mut self,
         run_instruction: Option<Run>,
         default_shell: TerminalAction,
+        tab_index: usize,
     ) -> Result<Option<(u32, bool, Option<RunCommand>, Result<i32>)>> {
         // terminal_id,
         // starts_held,
@@ -730,6 +737,7 @@ impl Pty {
                                 pane_id,
                                 exit_status,
                                 command,
+                                Some(tab_index),
                                 None,
                             ));
                         } else {
@@ -935,6 +943,7 @@ impl Pty {
                                 exit_status,
                                 command,
                                 None,
+                                None,
                             ));
                         } else {
                             let _ =
@@ -996,6 +1005,7 @@ fn send_command_not_found_to_screen(
     senders: ThreadSenders,
     terminal_id: u32,
     run_command: RunCommand,
+    tab_index: Option<usize>,
 ) -> Result<()> {
     let err_context = || format!("failed to send command_not_fount for terminal {terminal_id}");
     senders
@@ -1011,6 +1021,7 @@ fn send_command_not_found_to_screen(
             PaneId::Terminal(terminal_id),
             Some(2),
             run_command.clone(),
+            tab_index,
             None,
         ))
         .with_context(err_context)?;
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index e2cac0c91..eb0d359ac 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -175,7 +175,13 @@ pub enum ScreenInstruction {
     TogglePaneFrames,
     SetSelectable(PaneId, bool, usize),
     ClosePane(PaneId, Option<ClientId>),
-    HoldPane(PaneId, Option<i32>, RunCommand, Option<ClientId>), // Option<i32> is the exit status
+    HoldPane(
+        PaneId,
+        Option<i32>,
+        RunCommand,
+        Option<usize>,
+        Option<ClientId>,
+    ), // Option<i32> is the exit status, Option<usize> is the tab_index
     UpdatePaneName(Vec<u8>, ClientId),
     UndoRenamePane(ClientId),
     NewTab(
@@ -1914,10 +1920,10 @@ pub(crate) fn screen_thread_main(
                 screen.update_tabs()?;
                 screen.unblock_input()?;
             },
-            ScreenInstruction::HoldPane(id, exit_status, run_command, client_id) => {
+            ScreenInstruction::HoldPane(id, exit_status, run_command, tab_index, client_id) => {
                 let is_first_run = false;
-                match client_id {
-                    Some(client_id) => {
+                match (client_id, tab_index) {
+                    (Some(client_id), _) => {
                         active_tab!(screen, client_id, |tab: &mut Tab| tab.hold_pane(
                             id,
                             exit_status,
@@ -1925,7 +1931,14 @@ pub(crate) fn screen_thread_main(
                             run_command
                         ));
                     },
-                    None => {
+                    (_, Some(tab_index)) => {
+                        let tab = screen
+                            .tabs
+                            .get_mut(&tab_index)
+                            .context("couldn't find tab with index {tab_index}")?;
+                        tab.hold_pane(id, exit_status, is_first_run, run_command);
+                    },
+                    _ => {
                         for tab in screen.tabs.values_mut() {
                             if tab.get_all_pane_ids().contains(&id) {
                                 tab.hold_pane(id, exit_status, is_first_run, run_command);
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index 8a8133103..9e11fdfdf 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -89,6 +89,7 @@ type HoldForCommand = Option<RunCommand>;
 enum BufferedTabInstruction {
     SetPaneSelectable(PaneId, bool),
     HandlePtyBytes(u32, VteBytes),
+    HoldPane(PaneId, Option<i32>, bool, RunCommand), // Option<i32> is the exit status, bool is is_first_run
 }
 
 pub(crate) struct Tab {
@@ -737,6 +738,14 @@ impl Tab {
                 BufferedTabInstruction::HandlePtyBytes(terminal_id, bytes) => {
                     self.handle_pty_bytes(terminal_id, bytes)?;
                 },
+                BufferedTabInstruction::HoldPane(
+                    terminal_id,
+                    exit_status,
+                    is_first_run,
+                    run_command,
+                ) => {
+                    self.hold_pane(terminal_id, exit_status, is_first_run, run_command);
+                },
             }
         }
         Ok(())
@@ -2145,6 +2154,16 @@ impl Tab {
         is_first_run: bool,
         run_command: RunCommand,
     ) {
+        if self.is_pending {
+            self.pending_instructions
+                .push(BufferedTabInstruction::HoldPane(
+                    id,
+                    exit_status,
+                    is_first_run,
+                    run_command,
+                ));
+            return;
+        }
         if self.floating_panes.panes_contain(&id) {
             self.floating_panes
                 .hold_pane(id, exit_status, is_first_run, run_command);