From ac47ff43a384bd0f5f26903da6cc016c37c3d8af Mon Sep 17 00:00:00 2001
From: Aram Drevekenin <aram@poor.dev>
Date: Fri, 1 Dec 2023 19:10:13 +0100
Subject: [PATCH] fix(plugins): display errors properly (#2975)

* fix(plugins): format errors and make sure file not found errors appear

* style(fmt): rustfmt
---
 zellij-server/src/screen.rs                | 47 ++++++++++++++++++----
 zellij-server/src/ui/loading_indication.rs |  6 ++-
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index 409a1854..b11d09e3 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -2058,6 +2058,22 @@ impl Screen {
         }
         session_layout_metadata
     }
+    fn update_plugin_loading_stage(
+        &mut self,
+        pid: u32,
+        loading_indication: LoadingIndication,
+    ) -> bool {
+        let all_tabs = self.get_tabs_mut();
+        let mut found_plugin = false;
+        for tab in all_tabs.values_mut() {
+            if tab.has_plugin(pid) {
+                found_plugin = true;
+                tab.update_plugin_loading_stage(pid, loading_indication);
+                break;
+            }
+        }
+        found_plugin
+    }
 }
 
 // The box is here in order to make the
@@ -2117,6 +2133,7 @@ pub(crate) fn screen_thread_main(
     let mut pending_tab_switches: HashSet<(usize, ClientId)> = HashSet::new(); // usize is the
                                                                                // tab_index
 
+    let mut plugin_loading_message_cache = HashMap::new();
     loop {
         let (event, mut err_ctx) = screen
             .bus
@@ -2793,9 +2810,9 @@ pub(crate) fn screen_thread_main(
                 screen.apply_layout(
                     layout,
                     floating_panes_layout,
-                    new_pane_pids,
+                    new_pane_pids.clone(),
                     new_floating_pane_pids,
-                    new_plugin_ids,
+                    new_plugin_ids.clone(),
                     tab_index,
                     client_id,
                 )?;
@@ -2805,6 +2822,18 @@ pub(crate) fn screen_thread_main(
                         screen.go_to_tab(tab_index as usize, client_id)?;
                     }
                 }
+
+                for plugin_ids in new_plugin_ids.values() {
+                    for plugin_id in plugin_ids {
+                        if let Some(loading_indication) =
+                            plugin_loading_message_cache.remove(plugin_id)
+                        {
+                            screen.update_plugin_loading_stage(*plugin_id, loading_indication);
+                            screen.render()?;
+                        }
+                    }
+                }
+
                 screen.unblock_input()?;
                 screen.render()?;
             },
@@ -3307,16 +3336,18 @@ pub(crate) fn screen_thread_main(
                 } else {
                     log::error!("Tab index not found: {:?}", tab_index);
                 }
+                if let Some(loading_indication) = plugin_loading_message_cache.remove(&plugin_id) {
+                    screen.update_plugin_loading_stage(plugin_id, loading_indication);
+                    screen.render()?;
+                }
                 screen.log_and_report_session_state()?;
                 screen.unblock_input()?;
             },
             ScreenInstruction::UpdatePluginLoadingStage(pid, loading_indication) => {
-                let all_tabs = screen.get_tabs_mut();
-                for tab in all_tabs.values_mut() {
-                    if tab.has_plugin(pid) {
-                        tab.update_plugin_loading_stage(pid, loading_indication);
-                        break;
-                    }
+                let found_plugin =
+                    screen.update_plugin_loading_stage(pid, loading_indication.clone());
+                if !found_plugin {
+                    plugin_loading_message_cache.insert(pid, loading_indication);
                 }
                 screen.render()?;
             },
diff --git a/zellij-server/src/ui/loading_indication.rs b/zellij-server/src/ui/loading_indication.rs
index 1026c8ae..1c77d039 100644
--- a/zellij-server/src/ui/loading_indication.rs
+++ b/zellij-server/src/ui/loading_indication.rs
@@ -272,7 +272,11 @@ impl Display for LoadingIndication {
             None => {},
         }
         if let Some(error_text) = &self.error {
-            stringified.push_str(&format!("\n\r{} {error_text}", red.bold().paint("ERROR:")));
+            stringified.push_str(&format!(
+                "\n\r{} {}",
+                red.bold().paint("ERROR: "),
+                error_text.replace('\n', "\n\r")
+            ));
             // we add this additional line explicitly to make it easier to realize when something
             // is wrong in very small plugins (eg. the tab-bar and status-bar)
             stringified.push_str(&format!(