improve error handling in ui module (#1870)

* improve error handling in ui module

* resolve problems in the review

5 files changed, 85 insertions, 40 deletions

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)) = self
+            .pane
+            .render(None)
+            .context("failed to render pane contents to multiple clients")?
         {
             let clients: Vec<ClientId> = clients.collect();
             self.output.add_character_chunks_to_multiple_clients(
@@ -71,10 +74,13 @@ impl<'a> PaneContentsAndUi<'a> {
                 );
             }
         }
+        Ok(())
     }
-    pub fn render_pane_contents_for_client(&mut self, client_id: ClientId) {
-        if let Some((character_chunks, raw_vte_output, sixel_image_chunks)) =
-            self.pane.render(Some(client_id)).unwrap()
+    pub fn render_pane_contents_for_client(&mut self, client_id: ClientId) -> Result<()> {
+        if let Some((character_chunks, raw_vte_output, sixel_image_chunks)) = self
+            .pane
+            .render(Some(client_id))
+            .with_context(|| format!("failed to render pane contents for client {client_id}"))?
         {
             self.output
                 .add_character_chunks_to_client(client_id, character_chunks, self.z_index);
@@ -95,8 +101,9 @@ impl<'a> PaneContentsAndUi<'a> {
                 );
             }
         }
+        Ok(())
     }
-    pub fn render_fake_cursor_if_needed(&mut self, client_id: ClientId) {
+    pub fn render_fake_cursor_if_needed(&mut self, client_id: ClientId) -> Result<()> {
         let pane_focused_for_client_id = self.focused_clients.contains(&client_id);
         let pane_focused_for_different_client = self
             .focused_clients
@@ -109,7 +116,9 @@ impl<'a> PaneContentsAndUi<'a> {
                 .focused_clients
                 .iter()
                 .find(|&&c_id| c_id != client_id)
-                .unwrap();
+                .with_context(|| {
+                    format!("failed to render fake cursor if needed for client {client_id}")
+                })?;
             if let Some(colors) = client_id_to_colors(*fake_cursor_client_id, self.style.colors) {
                 if let Some(vte_output) = self.pane.render_fake_cursor(colors.0, colors.1) {
                     self.output.add_post_vte_instruction_to_client(
@@ -124,6 +133,7 @@ impl<'a> PaneContentsAndUi<'a> {
                 }
             }
         }
+        Ok(())
     }
     pub fn render_terminal_title_if_needed(&mut self, client_id: ClientId, client_mode: InputMode) {
         if !self.focused_clients.contains(&client_id) {
@@ -138,7 +148,8 @@ impl<'a> PaneContentsAndUi<'a> {
         client_id: ClientId,
         client_mode: InputMode,
         session_is_mirrored: bool,
-    ) {
+    ) -> Result<()> {
+        let err_context = || format!("failed to render pane frame for client {client_id}");
         let pane_focused_for_client_id = self.focused_clients.contains(&client_id);
         let other_focused_clients: Vec<ClientId> = self
             .focused_clients
@@ -152,7 +163,7 @@ impl<'a> PaneContentsAndUi<'a> {
         let focused_client = if pane_focused_for_client_id {
             Some(client_id)
         } else if pane_focused_for_differet_client {
-            Some(*other_focused_clients.first().unwrap())
+            Some(*other_focused_clients.first().with_context(err_context)?)
         } else {
             None
         };
@@ -175,8 +186,10 @@ impl<'a> PaneContentsAndUi<'a> {
                 other_cursors_exist_in_session: self.multiple_users_exist_in_session,
             }
         };
-        if let Some((frame_terminal_characters, vte_output)) =
-            self.pane.render_frame(client_id, frame_params, client_mode)
+        if let Some((frame_terminal_characters, vte_output)) = self
+            .pane
+            .render_frame(client_id, frame_params, client_mode)
+            .with_context(err_context)?
         {
             self.output.add_character_chunks_to_client(
                 client_id,
@@ -188,6 +201,7 @@ impl<'a> PaneContentsAndUi<'a> {
                     .add_post_vte_instruction_to_client(client_id, &vte_output);
             }
         }
+        Ok(())
     }
     pub fn render_pane_boundaries(
         &self,