Fix unwraps

Author: lionel-

crates/ark/src/console/console_comm.rs

@@ -51,8 +51,7 @@ impl Console {
     }
 
     pub(super) fn comm_handle_close(&self, comm_id: &str) {
-        if self.is_ui_comm(comm_id) {
-            let ui = self.take_ui_comm().unwrap();
+        if let Some(ui) = self.take_ui_comm_if(comm_id) {
             ui.handler.into_inner().handle_close(&ui.ctx);
             return;
         }
@@ -189,6 +188,17 @@ impl Console {
         self.ui_comm.borrow_mut().take()
     }
 
+    /// Take the UI comm only if its `comm_id` matches. Checks and takes
+    /// in a single `borrow_mut()` so there is no TOCTOU gap.
+    fn take_ui_comm_if(&self, comm_id: &str) -> Option<ConsoleComm> {
+        let mut guard = self.ui_comm.borrow_mut();
+        if guard.as_ref().is_some_and(|ui| ui.comm_id == comm_id) {
+            guard.take()
+        } else {
+            None
+        }
+    }
+
     fn set_ui_comm(&self, ui: ConsoleComm) {
         *self.ui_comm.borrow_mut() = Some(ui);
     }

crates/ark/src/plots/graphics_device.rs

@@ -704,7 +704,9 @@ impl DeviceContext {
             display_id: id.to_string(),
             data: None,
         };
-        let transient = serde_json::to_value(transient).unwrap();
+        let Some(transient) = serde_json::to_value(transient).log_err() else {
+            return;
+        };
 
         log::info!("Sending display data to IOPub.");
 
@@ -778,7 +780,10 @@ impl DeviceContext {
             },
         };
 
-        let value = serde_json::to_value(PlotFrontendEvent::Update(update_params)).unwrap();
+        let Some(value) = serde_json::to_value(PlotFrontendEvent::Update(update_params)).log_err()
+        else {
+            return;
+        };
 
         let outgoing_tx = CommOutgoingTx::new(comm_id, self.iopub_tx.clone());
         outgoing_tx
@@ -841,7 +846,7 @@ impl DeviceContext {
         });
 
         let mut map = serde_json::Map::new();
-        map.insert("image/png".to_string(), serde_json::to_value(data).unwrap());
+        map.insert("image/png".to_string(), serde_json::to_value(data)?);
 
         Ok(serde_json::Value::Object(map))
     }