WIP: handle update_display_data

Author: 3coins

jupyter_server_documents/kernels/kernel_client.py

@@ -223,6 +223,7 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
         parent_msg_data = self.message_cache.get(parent_msg_id)
         cell_id = parent_msg_data.get('cell_id')
 
+        self.log.info(f"msg_type: {dmsg["msg_type"]}")
         # Handle different message types using pattern matching
         match dmsg["msg_type"]:
             case "kernel_info_reply":
@@ -265,6 +266,7 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
                 if cell_id:
                     # Extract execution count and update each collaborative room's notebook
                     content = self.session.unpack(dmsg["content"])
+                    self.log.info(f"dsmg content: {content}")
                     execution_count = content["execution_count"]
                     for yroom in self._yrooms:
                         notebook = await yroom.get_jupyter_ydoc()
@@ -273,12 +275,13 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
                             target_cell["execution_count"] = execution_count
                             break
 
-            case "stream" | "display_data" | "execute_result" | "error":
+            case "stream" | "display_data" | "execute_result" | "error" | "update_display_data":
                 if cell_id: 
                     # Process specific output messages through an optional processor
                     if self.output_processor and cell_id:
                         cell_id = parent_msg_data.get('cell_id')
                         content = self.session.unpack(dmsg["content"])
+                        self.log.info(f"dsmg content: {content}")
                         dmsg = self.output_processor.process_output(dmsg['msg_type'], cell_id, content)
                         # Suppress forwarding of processed messages by returning None
                         return None

jupyter_server_documents/outputs/manager.py

@@ -14,6 +14,7 @@
 class OutputsManager(LoggingConfigurable):
     _last_output_index = Dict(default_value={})
     _stream_count = Dict(default_value={})
+    _output_index_by_display_id = Dict(default_value={})
 
     outputs_path = Instance(PurePath, help="The local runtime dir")
     stream_limit = Int(default_value=10, config=True, allow_none=True)
@@ -34,6 +35,10 @@ def _build_path(self, file_id, cell_id=None, output_index=None):
             path = path / f"{output_index}.output"
         return path
 
+    def get_output_index(self, display_id: str):
+        """Returns output index for a cell by display_id"""
+        return self._output_index_by_display_id.get(display_id)
+
     def get_output(self, file_id, cell_id, output_index):
         """Get an output by file_id, cell_id, and output_index."""
         path = self._build_path(file_id, cell_id, output_index)
@@ -102,23 +107,34 @@ def get_stream(self, file_id, cell_id):
         with open(path, "r", encoding="utf-8") as f:
             output = f.read()
         return output
-
-    def write(self, file_id, cell_id, output):
+    
+    def write(self, file_id, cell_id, output, display_id=None):
         """Write a new output for file_id and cell_id.
 
         Returns a placeholder output (pycrdt.Map) or None if no placeholder
         output should be written to the ydoc.
         """
-        placeholder = self.write_output(file_id, cell_id, output)
+        placeholder = self.write_output(file_id, cell_id, output, display_id)
         if output["output_type"] == "stream" and self.stream_limit is not None:
             placeholder = self.write_stream(file_id, cell_id, output, placeholder)
         return placeholder
 
-    def write_output(self, file_id, cell_id, output):
+    def write_output(self, file_id, cell_id, output, display_id=None):
+        print(f"display_id:{display_id} passed into write_output")
         self._ensure_path(file_id, cell_id)
         last_index = self._last_output_index.get(cell_id, -1)
-        index = last_index + 1
-        self._last_output_index[cell_id] = index
+        if display_id:
+            index = self._output_index_by_display_id.get(display_id)
+            print(f"Found index: {index} for display_id {display_id}")
+            if index is None:
+                index = last_index + 1
+                self._last_output_index[cell_id] = index
+                self._output_index_by_display_id[display_id] = index
+                print(f"Generated new index: {index} for display_id: {display_id}")
+        else:
+            index = last_index + 1
+            self._last_output_index[cell_id] = index
+        
         path = self._build_path(file_id, cell_id, index)
         data = json.dumps(output, ensure_ascii=False)
         with open(path, "w", encoding="utf-8") as f:
@@ -165,6 +181,10 @@ def clear(self, file_id, cell_id=None):
                 del self._stream_count[cell_id]
             except KeyError:
                 pass
+            try:
+                del self._last_output_index[cell_id]
+            except KeyError:
+                pass
             path = self._build_path(file_id, cell_id)
         try:
             shutil.rmtree(path)

jupyter_server_documents/outputs/output_processor.py

@@ -106,10 +106,13 @@ async def output_task(self, msg_type, cell_id, content):
             return
         self._file_id = file_id
 
+        display_id = content.get("transient", {}).get("display_id")
+        self.log.info(f"display_id: {display_id} found in content")
+        output_index = self.outputs_manager.get_output_index(display_id) if display_id else None
         # Convert from the message spec to the nbformat output structure
-        if self.use_outputs_service:
+        if self.use_outputs_service: 
             output = self.transform_output(msg_type, content, ydoc=False)
-            output = self.outputs_manager.write(file_id, cell_id, output)
+            output = self.outputs_manager.write(file_id, cell_id, output, display_id)
         else:
             output = self.transform_output(msg_type, content, ydoc=True)
 
@@ -125,8 +128,11 @@ async def output_task(self, msg_type, cell_id, content):
         # Write the outputs to the ydoc cell.
         _, target_cell = notebook.find_cell(cell_id)
         if target_cell is not None and output is not None:
-            target_cell["outputs"].append(output)
-            self.log.info(f"Write output to ydoc: {path} {cell_id} {output}")
+            if output_index:
+                target_cell["outputs"][output_index] = output
+            else:
+                target_cell["outputs"].append(output)
+            self.log.info(f"Wrote output to ydoc: {path} {cell_id} {output}")
 
 
     def transform_output(self, msg_type, content, ydoc=False):
@@ -141,7 +147,7 @@ def transform_output(self, msg_type, content, ydoc=False):
                 "text": content["text"],
                 "name": content["name"]
             })
-        elif msg_type == "display_data":
+        elif msg_type == "display_data" or msg_type == "update_display_data":
             output = factory({
                 "output_type": "display_data",
                 "data": content["data"],

src/notebook-factory/notebook-factory.ts

@@ -38,6 +38,7 @@ const DIRTY_CLASS = 'jp-mod-dirty';
   }
 
   if (change.outputsChange) {
+    console.trace(change.outputsChange);
     globalModelDBMutex(() => {
       let retain = 0;
       for (const outputsChange of change.outputsChange!) {