zellij-server/src/screen: improve error handling (#1770)

* zellij-server/src/screen: improve error handling

* cleanup context vs with_context usage

* update error handling docs

* zellij-server/src/screen.rs: fix formatting

2 files changed, 96 insertions, 57 deletions

diff --git a/docs/ERROR_HANDLING.md b/docs/ERROR_HANDLING.md
index 5479ebf4a..02ba2983c 100644
--- a/docs/ERROR_HANDLING.md
+++ b/docs/ERROR_HANDLING.md
@@ -26,7 +26,6 @@ If you have an interest in this, don't hesitate to get in touch with us and
 refer to the [tracking issue][tracking_issue] to see what has already been
 done.
 
-
 # Error handling facilities
 
 You get access to all the relevant functions and traits mentioned in the
@@ -43,7 +42,6 @@ Panics are generally handled via the `Panic` error type and the
 [`handle_panic`][handle_panic] panic handler function. The fancy formatting
 is performed by the [`miette`][miette] crate.
 
-
 ## Propagating errors
 
 We use the [`anyhow`][anyhow] crate to propagate errors up the call stack. At
@@ -53,7 +51,7 @@ of providing [`context`][context] about where (i.e. under which circumstances)
 an error happened.
 
 A critical requirement for propagating errors is that all functions involved
-must return the [`Result`][Result] type. This allows convenient error handling
+must return the [`Result`][result] type. This allows convenient error handling
 with the `?` operator.
 
 At some point you will likely stop propagating errors and decide what to do
@@ -64,7 +62,6 @@ with the error. Generally you can:
     1. Terminate program execution (See [`fatal`][fatal]), or
     2. Continue program execution (See [`non_fatal`][non_fatal])
 
-
 ## Handling errors
 
 Ideally, when the program encounters an error it will try to recover as good as
@@ -76,13 +73,13 @@ Recovery usually isn't an option if an operation has changed the internal state
 (i.e. the value or content of specific variables) of objects in the code. In
 this case, if an error is encountered, it is best to declare the program state
 corrupted and terminate the whole application. This can be done by `unwrap`ing
-on the [`Result`][Result] type. Always try to propagate the error as good as
+on the [`Result`][result] type. Always try to propagate the error as good as
 you can and attach meaningful context before `unwrap`ing. This gives the user
 an idea what went wrong and can also help developers in quickly identifying
 which parts of the code to debug if necessary.
 
 When you encounter such a fatal error and cannot propagate it further up (e.g.
-because the current function cannot be changed to return a [`Result`][Result],
+because the current function cannot be changed to return a [`Result`][result],
 or because it is the "root" function of a program thread), use the
 [`fatal`][fatal] function to panic the application. It will attach some small
 context to the error and finally `unwrap` it. Using this function over the
@@ -95,7 +92,6 @@ to handle it. Instead of `panic`ing the application, the error is written to
 the application log and execution continues. Please use this sparingly, as an
 error usually calls for actions to be taken rather than ignoring it.
 
-
 # Examples of applied error handling
 
 You can have a look at the commit that introduced error handling to the
@@ -103,7 +99,6 @@ You can have a look at the commit that introduced error handling to the
 `zellij-server/src/screen.rs`). We'll use this to demonstrate a few things in
 the following text.
 
-
 ## Converting a function to return a `Result` type
 
 Here's an example of the `Screen::render` function as it looked before:
@@ -172,7 +167,7 @@ to this:
     pub fn resize_to_screen(&mut self, new_screen_size: Size) -> Result<()> {
         // ...
         self.render()
-            .context("failed to resize to screen size: {new_screen_size:#?}")
+            .with_context(|| format!("failed to resize to screen size: {new_screen_size:#?}"))
     }
 ```
 
@@ -183,30 +178,58 @@ time to do something about the error.
 In general, any function calling `unwrap` or `expect` is a good candidate to be
 rewritten to return a `Result` type instead.
 
-
 ## Attaching context
 
 [Anyhow][anyhow]s [`Context`][context] trait gives us two methods to attach
-context to an error: `context` and `with_context`. In functions where `Result`s
-can be returned in multiple places, it is preferable to use `with_context`, as
-it saves you from duplicating the same context message in multiple places in
-the code. This was shown in the `render` method above, but the call to the
-other function returning a `Result` wasn't shown:
+context to an error: `context` and `with_context`. You should use `context`
+if the message contains only a static text and `with_context` if you need
+additional formatting:
+
+```rust
+    fn move_clients_between_tabs(
+        &mut self,
+        source_tab_index: usize,
+        destination_tab_index: usize,
+        clients_to_move: Option<Vec<ClientId>>,
+    ) -> Result<()> {
+        // ...
+        if let Some(client_mode_info_in_source_tab) = drained_clients {
+            let destination_tab = self.get_indexed_tab_mut(destination_tab_index)
+                .context("failed to get destination tab by index")
+                .with_context(|| format!("failed to move clients from tab {source_tab_index} to tab {destination_tab_index}"))?;
+            // ...
+        }
+        Ok(())
+    }
+```
+
+Feel free to move context string/closure to a variable to avoid copy-pasting:
 
 ```rust
     pub fn render(&mut self) -> Result<()> {
-        let err_context = || "failed to render screen".to_string();
+        let err_context = "failed to render screen";
         // ...
 
         for tab_index in tabs_to_close {
             // ...
             self.close_tab_at_index(tab_index)
-                .with_context(err_context)?;
+                .context(err_context)?;
         }
         // ...
         self.bus
             .senders
             .send_to_server(ServerInstruction::Render(Some(serialized_output)))
+            .context(err_context)
+    }
+    // ...
+    pub fn close_tab(&mut self, client_id: ClientId) -> Result<()> {
+        let err_context = || format!("failed to close tab for client {client_id:?}");
+
+        let active_tab_index = *self
+            .active_tab_indices
+            .get(&client_id)
+            .with_context(err_context)?;
+        self.close_tab_at_index(active_tab_index)
             .with_context(err_context)
     }
 ```
@@ -214,7 +237,6 @@ other function returning a `Result` wasn't shown:
 When there is only a single `Result` to be returned from your function, use
 `context` as shown above for `resize_to_screen`.
 
-
 ## Choosing helpful context messages
 
 When attaching context to an error, usually you want to express what you were
@@ -248,7 +270,6 @@ give us a [`NotFound`][2] error), so we don't have to repeat that.
 In case of doubt, look at the name of the function you're currently working in
 and write a context message somehow mentioning this.
 
-
 ## Terminating execution
 
 We want to propagate errors as far up as we can. This way, every function along
@@ -300,13 +321,12 @@ the application (i.e. crash zellij). Since we made sure to attach context
 messages to the errors on their way up, we will see these messages in the
 resulting output!
 
-
 [tracking_issue]: https://github.com/zellij-org/zellij/issues/1753
 [handle_panic]: https://docs.rs/zellij-utils/latest/zellij_utils/errors/fn.handle_panic.html
 [miette]: https://crates.io/crates/miette
 [anyhow]: https://crates.io/crates/anyhow
 [context]: https://docs.rs/anyhow/latest/anyhow/trait.Context.html
-[Result]: https://doc.rust-lang.org/stable/std/result/enum.Result.html
+[result]: https://doc.rust-lang.org/stable/std/result/enum.Result.html
 [fatal]: https://docs.rs/zellij-utils/latest/zellij_utils/errors/trait.FatalError.html#tymethod.fatal
 [non_fatal]: https://docs.rs/zellij-utils/latest/zellij_utils/errors/trait.FatalError.html#tymethod.non_fatal
 [1]: https://github.com/zellij-org/zellij/commit/99e2bef8c68bd166cf89e90c8ffe8c02272ab4d3
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index b302db4be..ef4b4f609 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -394,18 +394,28 @@ impl Screen {
         source_tab_index: usize,
         destination_tab_index: usize,
         clients_to_move: Option<Vec<ClientId>>,
-    ) {
+    ) -> Result<()> {
+        let err_context = || {
+            format!(
+                "failed to move clients from tab {source_tab_index} to tab {destination_tab_index}"
+            )
+        };
+
         // None ==> move all clients
         let drained_clients = self
             .get_indexed_tab_mut(source_tab_index)
             .map(|t| t.drain_connected_clients(clients_to_move));
         if let Some(client_mode_info_in_source_tab) = drained_clients {
-            let destination_tab = self.get_indexed_tab_mut(destination_tab_index).unwrap();
+            let destination_tab = self
+                .get_indexed_tab_mut(destination_tab_index)
+                .context("failed to get destination tab by index")
+                .with_context(err_context)?;
             destination_tab.add_multiple_clients(client_mode_info_in_source_tab);
             destination_tab.update_input_modes();
             destination_tab.set_force_render();
             destination_tab.visible(true);
         }
+        Ok(())
     }
     fn update_client_tab_focus(&mut self, client_id: ClientId, new_tab_index: usize) {
         match self.active_tab_indices.remove(&client_id) {
@@ -439,7 +449,8 @@ impl Screen {
                 let current_tab_index = current_tab.index;
                 let new_tab_index = new_tab.index;
                 if self.session_is_mirrored {
-                    self.move_clients_between_tabs(current_tab_index, new_tab_index, None);
+                    self.move_clients_between_tabs(current_tab_index, new_tab_index, None)
+                        .with_context(err_context)?;
                     let all_connected_clients: Vec<ClientId> =
                         self.connected_clients.borrow().iter().copied().collect();
                     for client_id in all_connected_clients {
@@ -450,7 +461,8 @@ impl Screen {
                         current_tab_index,
                         new_tab_index,
                         Some(vec![client_id]),
-                    );
+                    )
+                    .with_context(err_context)?;
                     self.update_client_tab_focus(client_id, new_tab_index);
                 }
 
@@ -559,7 +571,7 @@ impl Screen {
             tab.set_force_render();
         }
         self.render()
-            .context("failed to resize to screen size: {new_screen_size:#?}")
+            .with_context(|| format!("failed to resize to screen size: {new_screen_size:#?}"))
     }
 
     pub fn update_pixel_dimensions(&mut self, pixel_dimensions: PixelDimensions) {
@@ -604,7 +616,7 @@ impl Screen {
 
     /// Renders this [`Screen`], which amounts to rendering its active [`Tab`].
     pub fn render(&mut self) -> Result<()> {
-        let err_context = || "failed to render screen".to_string();
+        let err_context = "failed to render screen";
 
         let mut output = Output::new(
             self.sixel_image_store.clone(),
@@ -622,14 +634,13 @@ impl Screen {
             }
         }
         for tab_index in tabs_to_close {
-            self.close_tab_at_index(tab_index)
-                .with_context(err_context)?;
+            self.close_tab_at_index(tab_index).context(err_context)?;
         }
         let serialized_output = output.serialize();
         self.bus
             .senders
             .send_to_server(ServerInstruction::Render(Some(serialized_output)))
-            .with_context(err_context)
+            .context(err_context)
     }
 
     /// Returns a mutable reference to this [`Screen`]'s tabs.
@@ -689,7 +700,7 @@ impl Screen {
         new_pids: Vec<RawFd>,
         client_id: ClientId,
     ) -> Result<()> {
-        let err_context = || format!("failed to create new tab for client {client_id:?}",);
+        let err_context = || format!("failed to create new tab for client {client_id:?}");
 
         let tab_index = self.get_new_tab_index();
         let position = self.tabs.len();
@@ -703,7 +714,7 @@ impl Screen {
             self.bus
                 .os_input
                 .as_ref()
-                .with_context(|| format!("failed to create new tab for client {client_id:?}"))?
+                .with_context(err_context)?
                 .clone(),
             self.bus.senders.clone(),
             self.max_panes,
@@ -795,7 +806,7 @@ impl Screen {
         }
         self.connected_clients.borrow_mut().remove(&client_id);
         self.update_tabs()
-            .with_context(|| format!("failed to remote client {client_id:?}"))
+            .with_context(|| format!("failed to remove client {client_id:?}"))
     }
 
     pub fn update_tabs(&self) -> Result<()> {
@@ -843,7 +854,8 @@ impl Screen {
     }
 
     pub fn update_active_tab_name(&mut self, buf: Vec<u8>, client_id: ClientId) -> Result<()> {
-        let s = str::from_utf8(&buf).unwrap();
+        let s = str::from_utf8(&buf)
+            .with_context(|| format!("failed to construct tab name from buf: {buf:?}"))?;
         if let Some(active_tab) = self.get_active_tab_mut(client_id) {
             match s {
                 "\0" => {
@@ -860,8 +872,9 @@ impl Screen {
                     }
                 },
             }
-            self.update_tabs()
-                .context("failed to update active tabs name for client id: {client_id:?}")
+            self.update_tabs().with_context(|| {
+                format!("failed to update active tabs name for client id: {client_id:?}")
+            })
         } else {
             log::error!("Active tab not found for client id: {client_id:?}");
             Ok(())
@@ -965,11 +978,12 @@ impl Screen {
         self.render()
     }
 
-    fn unblock_input(&self) {
+    fn unblock_input(&self) -> Result<()> {
         self.bus
             .senders
             .send_to_server(ServerInstruction::UnblockInputThread)
-            .unwrap();
+            .context("failed to send message to server")
+            .context("failed to unblock input")
     }
 }
 
@@ -1043,7 +1057,7 @@ pub(crate) fn screen_thread_main(
                         }
                     },
                 };
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.update_tabs()?;
 
                 screen.render()?;
@@ -1051,7 +1065,7 @@ pub(crate) fn screen_thread_main(
             ScreenInstruction::OpenInPlaceEditor(pid, client_id) => {
                 active_tab!(screen, client_id, |tab: &mut Tab| tab
                     .suppress_active_pane(pid, client_id));
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.update_tabs()?;
 
                 screen.render()?;
@@ -1059,28 +1073,28 @@ pub(crate) fn screen_thread_main(
             ScreenInstruction::TogglePaneEmbedOrFloating(client_id) => {
                 active_tab!(screen, client_id, |tab: &mut Tab| tab
                     .toggle_pane_embed_or_floating(client_id));
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.update_tabs()?; // update tabs so that the ui indication will be send to the plugins
                 screen.render()?;
             },
             ScreenInstruction::ToggleFloatingPanes(client_id, default_shell) => {
                 active_tab!(screen, client_id, |tab: &mut Tab| tab
                     .toggle_floating_panes(client_id, default_shell));
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.update_tabs()?; // update tabs so that the ui indication will be send to the plugins
                 screen.render()?;
             },
             ScreenInstruction::HorizontalSplit(pid, client_id) => {
                 active_tab!(screen, client_id, |tab: &mut Tab| tab
                     .horizontal_split(pid, client_id));
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.update_tabs()?;
                 screen.render()?;
             },
             ScreenInstruction::VerticalSplit(pid, client_id) => {
                 active_tab!(screen, client_id, |tab: &mut Tab| tab
                     .vertical_split(pid, client_id));
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.update_tabs()?;
                 screen.render()?;
             },
@@ -1143,7 +1157,7 @@ pub(crate) fn screen_thread_main(
             },
             ScreenInstruction::MoveFocusLeftOrPreviousTab(client_id) => {
                 screen.move_focus_left_or_previous_tab(client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::MoveFocusDown(client_id) => {
@@ -1158,7 +1172,7 @@ pub(crate) fn screen_thread_main(
             },
             ScreenInstruction::MoveFocusRightOrNextTab(client_id) => {
                 screen.move_focus_right_or_next_tab(client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::MoveFocusUp(client_id) => {
@@ -1311,22 +1325,22 @@ pub(crate) fn screen_thread_main(
             },
             ScreenInstruction::SwitchTabNext(client_id) => {
                 screen.switch_tab_next(client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::SwitchTabPrev(client_id) => {
                 screen.switch_tab_prev(client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::CloseTab(client_id) => {
                 screen.close_tab(client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::NewTab(layout, new_pane_pids, client_id) => {
                 screen.new_tab(layout, new_pane_pids, client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::GoToTab(tab_index, client_id) => {
@@ -1334,7 +1348,7 @@ pub(crate) fn screen_thread_main(
                     client_id.or_else(|| screen.active_tab_indices.keys().next().copied())
                 {
                     screen.go_to_tab(tab_index as usize, client_id)?;
-                    screen.unblock_input();
+                    screen.unblock_input()?;
                     screen.render()?;
                 }
             },
@@ -1433,7 +1447,7 @@ pub(crate) fn screen_thread_main(
             },
             ScreenInstruction::ToggleTab(client_id) => {
                 screen.toggle_tab(client_id)?;
-                screen.unblock_input();
+                screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::AddClient(client_id) => {
@@ -1448,25 +1462,30 @@ pub(crate) fn screen_thread_main(
             ScreenInstruction::AddOverlay(overlay, _client_id) => {
                 screen.get_active_overlays_mut().pop();
                 screen.get_active_overlays_mut().push(overlay);
-                screen.unblock_input();
+                screen.unblock_input()?;
             },
             ScreenInstruction::RemoveOverlay(_client_id) => {
                 screen.get_active_overlays_mut().pop();
                 screen.render()?;
-                screen.unblock_input();
+                screen.unblock_input()?;
             },
             ScreenInstruction::ConfirmPrompt(_client_id) => {
                 let overlay = screen.get_active_overlays_mut().pop();
                 let instruction = overlay.and_then(|o| o.prompt_confirm());
                 if let Some(instruction) = instruction {
-                    screen.bus.senders.send_to_server(*instruction).unwrap();
+                    screen
+                        .bus
+                        .senders
+                        .send_to_server(*instruction)
+                        .context("failed to send message to server")
+                        .context("failed to confirm prompt")?;
                 }
-                screen.unblock_input();
+                screen.unblock_input()?;
             },
             ScreenInstruction::DenyPrompt(_client_id) => {
                 screen.get_active_overlays_mut().pop();
                 screen.render()?;
-                screen.unblock_input();
+                screen.unblock_input()?;
             },
             ScreenInstruction::UpdateSearch(c, client_id) => {
                 active_tab!(screen, client_id, |tab: &mut Tab| tab