WIP: Working version of clear_output

Author: 3coins

jupyter_server_documents/kernels/kernel_client.py

@@ -105,7 +105,6 @@ def handle_incoming_message(self, channel_name: str, msg: list[bytes]):
         msg_id = header["msg_id"]                
         metadata = self.session.unpack(msg[2])
         cell_id = metadata.get("cellId")
-        
         # Clear output processor if this cell already has 
         # an existing request.
         if cell_id:
@@ -218,7 +217,7 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
         except Exception as e:
             self.log.error(f"Error deserializing message: {e}")
             raise
-
+        
         parent_msg_id = dmsg["parent_header"]["msg_id"]
         parent_msg_data = self.message_cache.get(parent_msg_id)
         cell_id = parent_msg_data.get('cell_id')
@@ -273,13 +272,14 @@ 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" | "update_display_data":
+            case "stream" | "display_data" | "execute_result" | "error" | "update_display_data" | "clear_output":
                 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"])
-                        dmsg = self.output_processor.process_output(dmsg['msg_type'], cell_id, content)
+                        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/output_processor.py

@@ -2,13 +2,14 @@
 
 from pycrdt import Map
 
-from traitlets import Unicode, Bool
+from traitlets import Unicode, Bool, Dict
 from traitlets.config import LoggingConfigurable
 from jupyter_server_documents.kernels.message_cache import KernelMessageCache
 
 class OutputProcessor(LoggingConfigurable):
 
     _file_id = Unicode(default_value=None, allow_none=True)
+    _clear_output_cells = Dict(default_value={})
 
     use_outputs_service = Bool(
         default_value=True,
@@ -46,6 +47,17 @@ def yroom_manager(self):
         """A shortcut for the jupyter server ydoc manager."""
         return self.settings["yroom_manager"]
 
+
+    async def get_notebook_ydoc(self, file_id):
+        room_id = f"json:notebook:{file_id}"
+        room = self.yroom_manager.get_room(room_id)
+        if room is None:
+            self.log.error(f"YRoom not found: {room_id}")
+            return
+        notebook = await room.get_jupyter_ydoc()
+        
+        return notebook
+
     async def clear_cell_outputs(self, cell_id, room_id):
         """Clears the outputs of a cell in ydoc"""
         room = self.yroom_manager.get_room(room_id)
@@ -62,11 +74,19 @@ async def clear_cell_outputs(self, cell_id, room_id):
 
     def clear(self, cell_id):
         """Clear all outputs for a given cell Id."""
+
+        task = None
+
         if self._file_id is not None:
             if self.use_outputs_service:
                 room_id = f"json:notebook:{self._file_id}"
-                asyncio.create_task(self.clear_cell_outputs(cell_id, room_id))
+                task = asyncio.create_task(self.clear_cell_outputs(cell_id, room_id))
                 self.outputs_manager.clear(file_id=self._file_id, cell_id=cell_id)
+            
+            # Remove any pending delayed clear for this cell
+            self._clear_output_cells.pop(cell_id, None)
+        
+        return task
 
     # Outgoing messages
     def process_output(self, msg_type: str, cell_id: str, content: dict):
@@ -87,6 +107,23 @@ def process_output(self, msg_type: str, cell_id: str, content: dict):
 
     async def output_task(self, msg_type, cell_id, content):
         """A coroutine to handle output messages."""
+
+        if msg_type == "clear_output":
+            if not content.get("wait", False):
+                # Clear immediately
+                self.clear(cell_id)
+            else:
+                # Mark for delayed clear, overriding any previous pending clear
+                self._clear_output_cells[cell_id] = True
+
+            return
+
+        # Check for delayed clear before processing output
+        if cell_id in self._clear_output_cells and msg_type != "clear_output":
+            clear_task = self.clear(cell_id)
+            if clear_task is not None:
+                await clear_task
+
         try:
             # TODO: The session manager may have multiple notebooks connected to the kernel
             # but currently get_session only returns the first. We need to fix this and