improve error handling in ui module (#1870)

* improve error handling in ui module

* resolve problems in the review

12 files changed, 157 insertions, 77 deletions

diff --git a/zellij-server/src/panes/floating_panes/mod.rs b/zellij-server/src/panes/floating_panes/mod.rs
index 650a763a4..f686f40bb 100644
--- a/zellij-server/src/panes/floating_panes/mod.rs
+++ b/zellij-server/src/panes/floating_panes/mod.rs
@@ -17,6 +17,7 @@ use std::rc::Rc;
 use std::time::Instant;
 use zellij_utils::{
     data::{ModeInfo, Style},
+    errors::prelude::*,
     input::command::RunCommand,
     pane_size::{Offset, PaneGeom, Size, Viewport},
 };
@@ -234,7 +235,8 @@ impl FloatingPanes {
             resize_pty!(pane, os_api);
         }
     }
-    pub fn render(&mut self, output: &mut Output) {
+    pub fn render(&mut self, output: &mut Output) -> Result<()> {
+        let err_context = || "failed to render output";
         let connected_clients: Vec<ClientId> =
             { self.connected_clients.borrow().iter().copied().collect() };
         let mut floating_panes: Vec<_> = self.panes.iter_mut().collect();
@@ -242,8 +244,16 @@ impl FloatingPanes {
             self.z_indices
                 .iter()
                 .position(|id| id == *a_id)
-                .unwrap()
-                .cmp(&self.z_indices.iter().position(|id| id == *b_id).unwrap())
+                .with_context(err_context)
+                .fatal()
+                .cmp(
+                    &self
+                        .z_indices
+                        .iter()
+                        .position(|id| id == *b_id)
+                        .with_context(err_context)
+                        .fatal(),
+                )
         });
 
         for (z_index, (kind, pane)) in floating_panes.iter_mut().enumerate() {
@@ -266,23 +276,27 @@ impl FloatingPanes {
                     .get(client_id)
                     .unwrap_or(&self.default_mode_info)
                     .mode;
-                pane_contents_and_ui.render_pane_frame(
-                    *client_id,
-                    client_mode,
-                    self.session_is_mirrored,
-                );
+                pane_contents_and_ui
+                    .render_pane_frame(*client_id, client_mode, self.session_is_mirrored)
+                    .with_context(err_context)?;
                 if let PaneId::Plugin(..) = kind {
-                    pane_contents_and_ui.render_pane_contents_for_client(*client_id);
+                    pane_contents_and_ui
+                        .render_pane_contents_for_client(*client_id)
+                        .with_context(err_context)?;
                 }
                 // this is done for panes that don't have their own cursor (eg. panes of
                 // another user)
-                pane_contents_and_ui.render_fake_cursor_if_needed(*client_id);
+                pane_contents_and_ui
+                    .render_fake_cursor_if_needed(*client_id)
+                    .with_context(err_context)?;
             }
             if let PaneId::Terminal(..) = kind {
                 pane_contents_and_ui
-                    .render_pane_contents_to_multiple_clients(connected_clients.iter().copied());
+                    .render_pane_contents_to_multiple_clients(connected_clients.iter().copied())
+                    .with_context(err_context)?;
             }
         }
+        Ok(())
     }
     pub fn resize(&mut self, new_screen_size: Size) {
         let display_area = *self.display_area.borrow();
diff --git a/zellij-server/src/panes/plugin_pane.rs b/zellij-server/src/panes/plugin_pane.rs
index 5709d7fee..81610188f 100644
--- a/zellij-server/src/panes/plugin_pane.rs
+++ b/zellij-server/src/panes/plugin_pane.rs
@@ -231,10 +231,10 @@ impl Pane for PluginPane {
         _client_id: ClientId,
         frame_params: FrameParams,
         input_mode: InputMode,
-    ) -> Option<(Vec<CharacterChunk>, Option<String>)> {
+    ) -> Result<Option<(Vec<CharacterChunk>, Option<String>)>> {
         // FIXME: This is a hack that assumes all fixed-size panes are borderless. This
         // will eventually need fixing!
-        if self.frame && !(self.geom.rows.is_fixed() || self.geom.cols.is_fixed()) {
+        let res = if self.frame && !(self.geom.rows.is_fixed() || self.geom.cols.is_fixed()) {
             let pane_title = if self.pane_name.is_empty()
                 && input_mode == InputMode::RenamePane
                 && frame_params.is_main_client
@@ -251,10 +251,15 @@ impl Pane for PluginPane {
                 pane_title,
                 frame_params,
             );
-            Some(frame.render())
+            Some(
+                frame
+                    .render()
+                    .with_context(|| format!("failed to render frame for client {_client_id}"))?,
+            )
         } else {
             None
-        }
+        };
+        Ok(res)
     }
     fn render_fake_cursor(
         &mut self,
diff --git a/zellij-server/src/panes/terminal_pane.rs b/zellij-server/src/panes/terminal_pane.rs
index 8d0651ad2..631df0c84 100644
--- a/zellij-server/src/panes/terminal_pane.rs
+++ b/zellij-server/src/panes/terminal_pane.rs
@@ -347,7 +347,8 @@ impl Pane for TerminalPane {
         client_id: ClientId,
         frame_params: FrameParams,
         input_mode: InputMode,
-    ) -> Option<(Vec<CharacterChunk>, Option<String>)> {
+    ) -> Result<Option<(Vec<CharacterChunk>, Option<String>)>> {
+        let err_context = || format!("failed to render frame for client {client_id}");
         // TODO: remove the cursor stuff from here
         let pane_title = if self.pane_name.is_empty()
             && input_mode == InputMode::RenamePane
@@ -398,12 +399,12 @@ impl Pane for TerminalPane {
             frame.add_exit_status(exit_status.as_ref().copied());
         }
 
-        match self.frame.get(&client_id) {
+        let res = match self.frame.get(&client_id) {
             // TODO: use and_then or something?
             Some(last_frame) => {
                 if &frame != last_frame {
                     if !self.borderless {
-                        let frame_output = frame.render();
+                        let frame_output = frame.render().with_context(err_context)?;
                         self.frame.insert(client_id, frame);
                         Some(frame_output)
                     } else {
@@ -415,14 +416,15 @@ impl Pane for TerminalPane {
             },
             None => {
                 if !self.borderless {
-                    let frame_output = frame.render();
+                    let frame_output = frame.render().with_context(err_context)?;
                     self.frame.insert(client_id, frame);
                     Some(frame_output)
                 } else {
                     None
                 }
             },
-        }
+        };
+        Ok(res)
     }
     fn render_fake_cursor(
         &mut self,
diff --git a/zellij-server/src/panes/tiled_panes/mod.rs b/zellij-server/src/panes/tiled_panes/mod.rs
index af7769f3b..35f7a3b93 100644
--- a/zellij-server/src/panes/tiled_panes/mod.rs
+++ b/zellij-server/src/panes/tiled_panes/mod.rs
@@ -12,6 +12,7 @@ use std::cell::RefCell;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::rc::Rc;
 use std::time::Instant;
+use zellij_utils::errors::prelude::*;
 use zellij_utils::{
     data::{ModeInfo, Style},
     input::command::RunCommand,
@@ -378,7 +379,8 @@ impl TiledPanes {
     pub fn has_panes(&self) -> bool {
         !self.panes.is_empty()
     }
-    pub fn render(&mut self, output: &mut Output, floating_panes_are_visible: bool) {
+    pub fn render(&mut self, output: &mut Output, floating_panes_are_visible: bool) -> Result<()> {
+        let err_context = || "failed to render tiled panes";
         let connected_clients: Vec<ClientId> =
             { self.connected_clients.borrow().iter().copied().collect() };
         let multiple_users_exist_in_session = { self.connected_clients_in_app.borrow().len() > 1 };
@@ -409,15 +411,17 @@ impl TiledPanes {
                         .get(client_id)
                         .unwrap_or(&self.default_mode_info)
                         .mode;
+                    let err_context =
+                        || format!("failed to render tiled panes for client {client_id}");
                     if let PaneId::Plugin(..) = kind {
-                        pane_contents_and_ui.render_pane_contents_for_client(*client_id);
+                        pane_contents_and_ui
+                            .render_pane_contents_for_client(*client_id)
+                            .with_context(err_context)?;
                     }
                     if self.draw_pane_frames {
-                        pane_contents_and_ui.render_pane_frame(
-                            *client_id,
-                            client_mode,
-                            self.session_is_mirrored,
-                        );
+                        pane_contents_and_ui
+                            .render_pane_frame(*client_id, client_mode, self.session_is_mirrored)
+                            .with_context(err_context)?;
                     } else {
                         let boundaries = client_id_to_boundaries
                             .entry(*client_id)
@@ -432,20 +436,27 @@ impl TiledPanes {
                     pane_contents_and_ui.render_terminal_title_if_needed(*client_id, client_mode);
                     // this is done for panes that don't have their own cursor (eg. panes of
                     // another user)
-                    pane_contents_and_ui.render_fake_cursor_if_needed(*client_id);
+                    pane_contents_and_ui
+                        .render_fake_cursor_if_needed(*client_id)
+                        .with_context(err_context)?;
                 }
                 if let PaneId::Terminal(..) = kind {
-                    pane_contents_and_ui.render_pane_contents_to_multiple_clients(
-                        connected_clients.iter().copied(),
-                    );
+                    pane_contents_and_ui
+                        .render_pane_contents_to_multiple_clients(connected_clients.iter().copied())
+                        .with_context(err_context)?;
                 }
             }
         }
         // render boundaries if needed
         for (client_id, boundaries) in &mut client_id_to_boundaries {
             // TODO: add some conditional rendering here so this isn't rendered for every character
-            output.add_character_chunks_to_client(*client_id, boundaries.render(), None);
+            output.add_character_chunks_to_client(
+                *client_id,
+                boundaries.render().with_context(err_context)?,
+                None,
+            );
         }
+        Ok(())
     }
     pub fn get_panes(&self) -> impl Iterator<Item = (&PaneId, &Box<dyn Pane>)> {
         self.panes.iter()
diff --git a/zellij-server/src/pty.rs b/zellij-server/src/pty.rs
index 5e72f8282..03f87fd7a 100644
--- a/zellij-server/src/pty.rs
+++ b/zellij-server/src/pty.rs
@@ -452,7 +452,7 @@ impl Pty {
             .bus
             .os_input
             .as_mut()
-            .unwrap()
+            .ok_or_else(|| SpawnTerminalError::GenericSpawnError("os input is none"))?
             .spawn_terminal(terminal_action, quit_cb, self.default_editor.clone())?;
         let terminal_bytes = task::spawn({
             let err_context = || format!("failed to run async task for terminal {terminal_id}");
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index 804bdf91f..910f217b3 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -736,7 +736,7 @@ impl Screen {
         let overlay = self.overlay.clone();
         for (tab_index, tab) in &mut self.tabs {
             if tab.has_selectable_tiled_panes() {
-                let vte_overlay = overlay.generate_overlay(size);
+                let vte_overlay = overlay.generate_overlay(size).context(err_context)?;
                 tab.render(&mut output, Some(vte_overlay))
                     .context(err_context)?;
             } else {
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index c63cac89c..2f8f51780 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -148,7 +148,7 @@ pub trait Pane {
         client_id: ClientId,
         frame_params: FrameParams,
         input_mode: InputMode,
-    ) -> Option<(Vec<CharacterChunk>, Option<String>)>; // TODO: better
+    ) -> Result<Option<(Vec<CharacterChunk>, Option<String>)>>; // TODO: better
     fn render_fake_cursor(
         &mut self,
         cursor_color: PaletteColor,
@@ -1329,9 +1329,12 @@ impl Tab {
 
         self.hide_cursor_and_clear_display_as_needed(output);
         self.tiled_panes
-            .render(output, self.floating_panes.panes_are_visible());
+            .render(output, self.floating_panes.panes_are_visible())
+            .with_context(err_context)?;
         if self.floating_panes.panes_are_visible() && self.floating_panes.has_active_panes() {
-            self.floating_panes.render(output);
+            self.floating_panes
+                .render(output)
+                .with_context(err_context)?;
         }
 
         // FIXME: Once clients can be distinguished
diff --git a/zellij-server/src/ui/boundaries.rs b/zellij-server/src/ui/boundaries.rs
index 4cd6b1fd7..c0827c24e 100644
--- a/zellij-server/src/ui/boundaries.rs
+++ b/zellij-server/src/ui/boundaries.rs
@@ -5,6 +5,7 @@ use crate::panes::terminal_character::{TerminalCharacter, EMPTY_TERMINAL_CHARACT
 use crate::tab::Pane;
 use ansi_term::Colour::{Fixed, RGB};
 use std::collections::HashMap;
+use zellij_utils::errors::prelude::*;
 use zellij_utils::{data::PaletteColor, shared::colors};
 
 use std::fmt::{Display, Error, Formatter};
@@ -47,18 +48,29 @@ impl BoundarySymbol {
         self.color = color;
         *self
     }
-    pub fn as_terminal_character(&self) -> TerminalCharacter {
-        if self.invisible {
+    pub fn as_terminal_character(&self) -> Result<TerminalCharacter> {
+        let tc = if self.invisible {
             EMPTY_TERMINAL_CHARACTER
         } else {
-            let character = self.boundary_type.chars().next().unwrap();
+            let character = self
+                .boundary_type
+                .chars()
+                .next()
+                .context("no boundary symbols defined")
+                .with_context(|| {
+                    format!(
+                        "failed to convert boundary symbol {} into terminal character",
+                        self.boundary_type
+                    )
+                })?;
             TerminalCharacter {
                 character,
                 width: 1,
                 styles: RESET_STYLES
                     .foreground(self.color.map(|palette_color| palette_color.into())),
             }
-        }
+        };
+        Ok(tc)
     }
 }
 
@@ -528,16 +540,18 @@ impl Boundaries {
             }
         }
     }
-    pub fn render(&self) -> Vec<CharacterChunk> {
+    pub fn render(&self) -> Result<Vec<CharacterChunk>> {
         let mut character_chunks = vec![];
         for (coordinates, boundary_character) in &self.boundary_characters {
             character_chunks.push(CharacterChunk::new(
-                vec![boundary_character.as_terminal_character()],
+                vec![boundary_character
+                    .as_terminal_character()
+                    .context("failed to render as terminal character")?],
                 coordinates.x,
                 coordinates.y,
             ));
         }
-        character_chunks
+        Ok(character_chunks)
     }
     fn rect_right_boundary_is_before_screen_edge(&self, rect: &dyn Pane) -> bool {
         rect.x() + rect.cols() < self.viewport.cols
diff --git a/zellij-server/src/ui/overlay/mod.rs b/zellij-server/src/ui/overlay/mod.rs
index ff9ac716a..7e3f821b3 100644
--- a/zellij-server/src/ui/overlay/mod.rs
+++ b/zellij-server/src/ui/overlay/mod.rs
@@ -9,6 +9,7 @@
 pub mod prompt;
 
 use crate::ServerInstruction;
+use zellij_utils::errors::prelude::*;
 use zellij_utils::pane_size::Size;
 
 #[derive(Clone, Debug)]
@@ -19,7 +20,7 @@ pub struct Overlay {
 pub trait Overlayable {
     /// Generates vte_output that can be passed into
     /// the `render()` function
-    fn generate_overlay(&self, size: Size) -> String;
+    fn generate_overlay(&self, size: Size) -> Result<String>;
 }
 
 #[derive(Clone, Debug)]
@@ -28,9 +29,11 @@ pub enum OverlayType {
 }
 
 impl Overlayable for OverlayType {
-    fn generate_overlay(&self, size: Size) -> String {
+    fn generate_overlay(&self, size: Size) -> Result<String> {
         match &self {
-            OverlayType::Prompt(prompt) => prompt.generate_overlay(size),
+            OverlayType::Prompt(prompt) => prompt
+                .generate_overlay(size)
+                .context("failed to generate VTE output from overlay type"),
         }
     }
 }
@@ -44,15 +47,17 @@ pub struct OverlayWindow {
 }
 
 impl Overlayable for OverlayWindow {
-    fn generate_overlay(&self, size: Size) -> String {
+    fn generate_overlay(&self, size: Size) -> Result<String> {
         let mut output = String::new();
         //let clear_display = "\u{1b}[2J";
         //output.push_str(&clear_display);
         for overlay in &self.overlay_stack {
-            let vte_output = overlay.generate_overlay(size);
+            let vte_output = overlay
+                .generate_overlay(size)
+                .context("failed to generate VTE output from overlay window")?;
             output.push_str(&vte_output);
         }
-        output
+        Ok(output)
     }
 }
 
@@ -70,8 +75,10 @@ impl Overlay {
 }
 
 impl Overlayable for Overlay {
-    fn generate_overlay(&self, size: Size) -> String {
-        self.overlay_type.generate_overlay(size)
+    fn generate_overlay(&self, size: Size) -> Result<String> {
+        self.overlay_type
+            .generate_overlay(size)
+            .context("failed to generate VTE output from overlay")
     }
 }
 
diff --git a/zellij-server/src/ui/overlay/prompt.rs b/zellij-server/src/ui/overlay/prompt.rs
index dc5171f52..b6b46f40f 100644
--- a/zellij-server/src/ui/overlay/prompt.rs
+++ b/zellij-server/src/ui/overlay/prompt.rs
@@ -2,6 +2,7 @@ use zellij_utils::pane_size::Size;
 
 use super::{Overlay, OverlayType, Overlayable};
 use crate::{ClientId, ServerInstruction};
+use zellij_utils::errors::prelude::*;
 
 use std::fmt::Write;
 
@@ -33,7 +34,7 @@ impl Prompt {
 }
 
 impl Overlayable for Prompt {
-    fn generate_overlay(&self, size: Size) -> String {
+    fn generate_overlay(&self, size: Size) -> Result<String> {
         let mut output = String::new();
         let rows = size.rows;
         let mut vte_output = self.message.clone();
@@ -46,9 +47,9 @@ impl Overlayable for Prompt {
                 x + 1,
                 h,
             )
-            .unwrap();
+            .context("failed to generate VTE output from prompt")?;
         }
-        output
+        Ok(output)
     }
 }
 
diff --git a/zellij-server/src/ui/pane_boundaries_frame.rs b/zellij-server/src/ui/pane_boundaries_frame.rs
index 186d4d182..61df8470a 100644
--- a/zellij-server/src/ui/pane_boundaries_frame.rs
+++ b/zellij-server/src/ui/pane_boundaries_frame.rs
@@ -3,6 +3,7 @@ use crate::panes::{AnsiCode, CharacterStyles, TerminalCharacter, EMPTY_TERMINAL_
 use crate::ui::boundaries::boundary_type;
 use crate::ClientId;
 use zellij_utils::data::{client_id_to_colors, PaletteColor, Style};
+use zellij_utils::errors::prelude::*;
 use zellij_utils::pane_size::Viewport;
 
 use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
@@ -600,24 +601,26 @@ impl PaneFrame {
             _ => self.empty_title_line(),
         }
     }
-    fn render_title(&self) -> Vec<TerminalCharacter> {
+    fn render_title(&self) -> Result<Vec<TerminalCharacter>> {
         let total_title_length = self.geom.cols.saturating_sub(2); // 2 for the left and right corners
 
         self.render_title_middle(total_title_length)
             .map(|(middle, middle_length)| self.title_line_with_middle(middle, &middle_length))
             .or_else(|| Some(self.title_line_without_middle()))
-            .unwrap()
+            .with_context(|| format!("failed to render title '{}'", self.title))
     }
-    fn render_held_undertitle(&self) -> Vec<TerminalCharacter> {
+    fn render_held_undertitle(&self) -> Result<Vec<TerminalCharacter>> {
         let max_undertitle_length = self.geom.cols.saturating_sub(2); // 2 for the left and right corners
-        let exit_status = self.exit_status.unwrap(); // unwrap is safe because we only call this if
+        let exit_status = self
+            .exit_status
+            .with_context(|| format!("failed to render command pane status '{}'", self.title))?; // unwrap is safe because we only call this if
 
         let (mut first_part, first_part_len) = self.first_held_title_part_full(exit_status);
         let mut left_boundary =
             foreground_color(self.get_corner(boundary_type::BOTTOM_LEFT), self.color);
         let mut right_boundary =
             foreground_color(self.get_corner(boundary_type::BOTTOM_RIGHT), self.color);
-        if self.is_main_client {
+        let res = if self.is_main_client {
             let (mut second_part, second_part_len) = self.second_held_title_part_full();
             let full_text_len = first_part_len + second_part_len;
             if full_text_len <= max_undertitle_length {
@@ -665,14 +668,16 @@ impl PaneFrame {
             } else {
                 self.empty_undertitle(max_undertitle_length)
             }
-        }
+        };
+        Ok(res)
     }
-    pub fn render(&self) -> (Vec<CharacterChunk>, Option<String>) {
+    pub fn render(&self) -> Result<(Vec<CharacterChunk>, Option<String>)> {
+        let err_context = || "failed to render pane frame";
         let mut character_chunks = vec![];
         for row in 0..self.geom.rows {
             if row == 0 {
                 // top row
-                let title = self.render_title();
+                let title = self.render_title().with_context(err_context)?;
                 let x = self.geom.x;
                 let y = self.geom.y + row;
                 character_chunks.push(CharacterChunk::new(title, x, y));
@@ -681,7 +686,11 @@ impl PaneFrame {
                 if self.exit_status.is_some() {
                     let x = self.geom.x;
                     let y = self.geom.y + row;
-                    character_chunks.push(CharacterChunk::new(self.render_held_undertitle(), x, y));
+                    character_chunks.push(CharacterChunk::new(
+                        self.render_held_undertitle().with_context(err_context)?,
+                        x,
+                        y,
+                    ));
                 } else {
                     let mut bottom_row = vec![];
                     for col in 0..self.geom.cols {
@@ -716,7 +725,7 @@ impl PaneFrame {
                 character_chunks.push(CharacterChunk::new(boundary_character_right, x, y));
             }
         }
-        (character_chunks, None)
+        Ok((character_chunks, None))
     }
     fn first_held_title_part_full(
         &self,
diff --git a/zellij-server/src/ui/pane_contents_and_ui.rs b/zellij-server/src/ui/pane_contents_and_ui.rs
index 4db9feb4b..f45b14a23 100644
--- a/zellij-server/src/ui/pane_contents_and_ui.rs
+++ b/zellij-server/src/ui/pane_contents_and_ui.rs
@@ -8,6 +8,7 @@ use std::collections::HashMap;
 use zellij_utils::data::{
     client_id_to_colors, single_client_color, InputMode, PaletteColor, Style,
 };
+use zellij_utils::errors::prelude::*;
 pub struct PaneContentsAndUi<'a> {
     pane: &'a mut Box<dyn Pane>,
     output: &'a mut Output,
@@ -44,9 +45,11 @@ impl<'a> PaneContentsAndUi<'a> {
     pub fn render_pane_contents_to_multiple_clients(
         &mut self,
         clients: impl Iterator<Item = ClientId>,
-    ) {
-        if let Some((character_chunks, raw_vte_output, sixel_image_chunks)) =
-            self.pane.render(None).unwrap()
+    ) -> Result<()> {
+        if let Some((character_chunks, raw_vte_output, sixel_image_chunks))