write outputs to ydoc client-side for jupyter-collaboration

Author: andrii-i

src/execution/remote/mod.rs

@@ -179,6 +179,7 @@ impl ExecutionBackend for RemoteExecutor {
         let cell_idx = cell_index.context("cell_index required for remote execution")?;
         let ydoc = self.ydoc.as_mut().context("Y.js client not connected")?;
         let http = reqwest::Client::new();
+        let client_writes = !ydoc.server_writes_outputs();
 
         // 1. Fire execute request
         let msg_id = ws
@@ -187,6 +188,7 @@ impl ExecutionBackend for RemoteExecutor {
 
         // 2. Watch for changes on the ydoc for this cell
         let mut outputs: Vec<nbformat::v4::Output> = Vec::new();
+        let mut kernel_outputs: Vec<nbformat::v4::Output> = Vec::new();
         let mut fetched_urls: HashSet<String> = HashSet::new();
         let mut seen_indices: HashSet<usize> = HashSet::new();
         let mut idle_received = false;
@@ -227,28 +229,20 @@ impl ExecutionBackend for RemoteExecutor {
                 }
 
                 if idle_received {
-                    let has_error = outputs
-                        .iter()
-                        .any(|o| matches!(o, nbformat::v4::Output::Error(_)));
-                    let error_info = outputs.iter().find_map(|o| {
-                        if let nbformat::v4::Output::Error(err) = o {
-                            Some(ExecutionError {
-                                ename: err.ename.clone(),
-                                evalue: err.evalue.clone(),
-                                traceback: err.traceback.clone(),
-                            })
-                        } else {
-                            None
-                        }
-                    });
-                    return if has_error {
-                        Ok(ExecutionResult::error(outputs, ec, error_info.unwrap()))
-                    } else {
-                        Ok(ExecutionResult::success(outputs, ec))
-                    };
+                    return Self::build_result(outputs, ec);
                 }
             }
 
+            // When the server doesn't write outputs and kernel is done,
+            // write collected outputs to Y.js ourselves, sync, then let
+            // the read loop above pick them up on the next iteration.
+            if client_writes && idle_received && !ec_ready && !kernel_outputs.is_empty() {
+                ydoc.update_cell_outputs(cell_idx, kernel_outputs.clone())?;
+                ydoc.update_cell_execution_count(cell_idx, expected_ec)?;
+                ydoc.sync().await?;
+                continue;
+            }
+
             // 4. Wait for new messages
             if idle_received {
                 match tokio::time::timeout_at(deadline, ydoc.recv_update()).await {
@@ -273,7 +267,13 @@ impl ExecutionBackend for RemoteExecutor {
                                             idle_received = true;
                                         }
                                     }
-                                    _ => {}
+                                    _ => {
+                                        if client_writes {
+                                            if let Some(output) = Self::kernel_msg_to_output(&msg.content) {
+                                                kernel_outputs.push(output);
+                                            }
+                                        }
+                                    }
                                 }
                             }
                         }
@@ -285,6 +285,11 @@ impl ExecutionBackend for RemoteExecutor {
             }
         }
 
+        // Fallback: if we collected kernel outputs but never wrote them
+        if client_writes && !kernel_outputs.is_empty() {
+            return Self::build_result(kernel_outputs, expected_ec);
+        }
+
         let ec = ydoc
             .read_cell_outputs(cell_idx)
             .ok()
@@ -310,3 +315,70 @@ impl ExecutionBackend for RemoteExecutor {
         Ok(())
     }
 }
+
+impl RemoteExecutor {
+    fn build_result(
+        outputs: Vec<nbformat::v4::Output>,
+        ec: Option<i64>,
+    ) -> Result<ExecutionResult> {
+        let has_error = outputs
+            .iter()
+            .any(|o| matches!(o, nbformat::v4::Output::Error(_)));
+        let error_info = outputs.iter().find_map(|o| {
+            if let nbformat::v4::Output::Error(err) = o {
+                Some(ExecutionError {
+                    ename: err.ename.clone(),
+                    evalue: err.evalue.clone(),
+                    traceback: err.traceback.clone(),
+                })
+            } else {
+                None
+            }
+        });
+        if has_error {
+            Ok(ExecutionResult::error(outputs, ec, error_info.unwrap()))
+        } else {
+            Ok(ExecutionResult::success(outputs, ec))
+        }
+    }
+
+    fn kernel_msg_to_output(content: &JupyterMessageContent) -> Option<nbformat::v4::Output> {
+        match content {
+            JupyterMessageContent::StreamContent(stream) => {
+                let name = match stream.name {
+                    jupyter_protocol::Stdio::Stdout => "stdout".to_string(),
+                    jupyter_protocol::Stdio::Stderr => "stderr".to_string(),
+                };
+                Some(nbformat::v4::Output::Stream {
+                    name,
+                    text: nbformat::v4::MultilineString(stream.text.clone()),
+                })
+            }
+            JupyterMessageContent::ExecuteResult(result) => {
+                let json = serde_json::json!({
+                    "output_type": "execute_result",
+                    "execution_count": result.execution_count.value(),
+                    "data": result.data,
+                    "metadata": result.metadata
+                });
+                serde_json::from_value(json).ok()
+            }
+            JupyterMessageContent::DisplayData(display) => {
+                let json = serde_json::json!({
+                    "output_type": "display_data",
+                    "data": display.data,
+                    "metadata": display.metadata
+                });
+                serde_json::from_value(json).ok()
+            }
+            JupyterMessageContent::ErrorOutput(error) => {
+                Some(nbformat::v4::Output::Error(nbformat::v4::ErrorOutput {
+                    ename: error.ename.clone(),
+                    evalue: error.evalue.clone(),
+                    traceback: error.traceback.clone(),
+                }))
+            }
+            _ => None,
+        }
+    }
+}

src/execution/remote/output_conversion.rs

@@ -6,7 +6,6 @@ use std::collections::HashMap;
 use yrs::{Any, Array, ArrayPrelim, ArrayRef, Map, MapPrelim, TransactionMut};
 
 /// Convert an nbformat Output to a MapPrelim that can be inserted into the outputs array
-#[allow(dead_code)]
 pub fn output_to_map_prelim(output: &Output) -> MapPrelim {
     match output {
         Output::Stream { name, text } => MapPrelim::from([
@@ -92,7 +91,6 @@ fn json_to_any(value: &JsonValue) -> Any {
 }
 
 /// Update a cell's outputs in the Y.js document
-#[allow(dead_code)]
 pub fn update_cell_outputs(
     txn: &mut TransactionMut,
     cells_array: &ArrayRef,
@@ -136,7 +134,6 @@ pub fn update_cell_outputs(
 }
 
 /// Update a cell's execution_count in the Y.js document
-#[allow(dead_code)]
 pub fn update_cell_execution_count(
     txn: &mut TransactionMut,
     cells_array: &ArrayRef,

src/execution/remote/ydoc.rs

@@ -66,6 +66,8 @@ pub struct YDocClient {
     file_id: String,
     /// Track the document state when we last synced, so we only send changes
     last_state: StateVector,
+    /// Whether the server writes outputs to Y.js (JSD does, jupyter-collaboration doesn't)
+    server_writes_outputs: bool,
 }
 
 impl YDocClient {
@@ -89,6 +91,7 @@ impl YDocClient {
             ws: ws_stream,
             file_id,
             last_state: StateVector::default(),
+            server_writes_outputs: session_id.is_none(),
         };
 
         // Step 4: Perform Y.js sync handshake with timeout
@@ -345,8 +348,11 @@ impl YDocClient {
         Ok(())
     }
 
+    pub fn server_writes_outputs(&self) -> bool {
+        self.server_writes_outputs
+    }
+
     /// Update cell outputs in the Y.js document
-    #[allow(dead_code)]
     pub fn update_cell_outputs(&mut self, cell_index: usize, outputs: Vec<Output>) -> Result<()> {
         let cells_array: ArrayRef = self.doc.get_or_insert_array("cells");
         let mut txn = self.doc.transact_mut();
@@ -358,7 +364,6 @@ impl YDocClient {
     }
 
     /// Update cell execution_count in the Y.js document
-    #[allow(dead_code)]
     pub fn update_cell_execution_count(
         &mut self,
         cell_index: usize,