Handle update_display_data kernel message in notebooks (#104)

* Added a method to get all outputs for a cell.

* Updated handler to accommodate all outputs

* Updated to sort outputs by last modified time

* Updated handler

* lint

* Updated to return jsonl response

* Sorting outputs by output index, simplified length check

* Moved placeholder to a separate function, ruff format.

* Updated method name

* WIP: handle update_display_data

* Fixed outputs index check

* Removed debug statements

* Moved output index logic to a separate method, added tests

* Fixed failing test

* Removed stale test

* fmt

* Fixed missing code after rebase

* Clean up after rebase

* Added output index tracker

* Updated tests

* Updates tests

* Remove unused imports

* Added a state to track display_ids, added tests

* Removed index tracker

* Renamed cell state vars

* Moved statement close to logic

Author: 3coins

jupyter_server_documents/kernels/kernel_client.py

@@ -273,7 +273,7 @@ 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:

jupyter_server_documents/outputs/manager.py

@@ -13,6 +13,8 @@
 
 class OutputsManager(LoggingConfigurable):
     _last_output_index = Dict(default_value={})
+    _output_index_by_display_id = Dict(default_value={})
+    _display_ids_by_cell_id = Dict(default_value={})
     _stream_count = Dict(default_value={})
 
     outputs_path = Instance(PurePath, help="The local runtime dir")
@@ -33,6 +35,38 @@ def _build_path(self, file_id, cell_id=None, output_index=None):
         if output_index is not None:
             path = path / f"{output_index}.output"
         return path
+    
+    def _compute_output_index(self, cell_id, display_id=None):
+        """
+        Computes next output index for a cell.
+        
+        Args:
+            cell_id (str): The cell identifier
+            display_id (str, optional): A display identifier. Defaults to None.
+        
+        Returns:
+            int: The output index
+        """
+        last_index = self._last_output_index.get(cell_id, -1)
+        if display_id:
+            if cell_id not in self._display_ids_by_cell_id:
+                self._display_ids_by_cell_id[cell_id] = set([display_id])
+            else:
+                self._display_ids_by_cell_id[cell_id].add(display_id)
+            index = self._output_index_by_display_id.get(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
+        else:
+            index = last_index + 1
+            self._last_output_index[cell_id] = index
+        
+        return index
+
+    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."""
@@ -77,23 +111,21 @@ 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):
         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
+        index = self._compute_output_index(cell_id, display_id)
         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:
@@ -134,13 +166,15 @@ def clear(self, file_id, cell_id=None):
         """Clear the state of the manager."""
         if cell_id is None:
             self._stream_count = {}
-            path = self._build_path(file_id)
         else:
-            try:
-                del self._stream_count[cell_id]
-            except KeyError:
-                pass
-            path = self._build_path(file_id, cell_id)
+            self._stream_count.pop(cell_id, None)
+            self._last_output_index.pop(cell_id, None)
+            
+            display_ids = self._display_ids_by_cell_id.get(cell_id, [])
+            for display_id in display_ids:
+                self._output_index_by_display_id.pop(display_id, None)
+
+        path = self._build_path(file_id, cell_id)    
         try:
             shutil.rmtree(path)
         except FileNotFoundError:

jupyter_server_documents/outputs/output_processor.py

@@ -1,9 +1,8 @@
 import asyncio
-import json
 
 from pycrdt import Map
 
-from traitlets import Dict, Unicode, Bool, Instance
+from traitlets import Unicode, Bool
 from traitlets.config import LoggingConfigurable
 from jupyter_server_documents.kernels.message_cache import KernelMessageCache
 
@@ -92,9 +91,13 @@ async def output_task(self, msg_type, cell_id, content):
             # 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
             # find the notebook with the right cell_id.
-            kernel_session = await self.session_manager.get_session(kernel_id=self.parent.parent.kernel_id)
+            kernel_session = await self.session_manager.get_session(
+                kernel_id=self.parent.parent.kernel_id
+            )
         except Exception as e:
-            self.log.error(f"An exception occurred when processing output for cell {cell_id}")
+            self.log.error(
+                f"An exception occurred when processing output for cell {cell_id}"
+            )
             self.log.exception(e)
             return
         else:
@@ -106,10 +109,11 @@ async def output_task(self, msg_type, cell_id, content):
             return
         self._file_id = file_id
 
+        display_id = content.get("transient", {}).get("display_id")
         # 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 +129,12 @@ 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}")
+            output_index = self.outputs_manager.get_output_index(display_id) if display_id else None
+            if output_index is not None:
+                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 +149,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"],

jupyter_server_documents/tests/test_output_processor.py

@@ -22,32 +22,6 @@ def test_instantiation():
     op = OutputProcessor()
     assert isinstance(op, OutputProcessor)
 
-def test_incoming_message():
-    """Test incoming message processing."""
-    with TemporaryDirectory() as td:
-        op = TestOutputProcessor()
-        om = OutputsManager()
-        op.settings["outputs_manager"] = om
-        op.outputs_path = Path(td) / "outputs"
-        # Simulate the running of a cell
-        cell_id = str(uuid4())
-        msg_id, msg = create_incoming_message(cell_id)
-        op.process_incoming_message('shell', msg)
-        assert op.get_cell_id(msg_id) == cell_id
-        assert op.get_msg_id(cell_id) == msg_id
-        # Clear the cell_id
-        op.clear(cell_id)
-        assert op.get_cell_id(msg_id) is None
-        assert op.get_msg_id(cell_id) is None
-        # Simulate the running of a cell
-        cell_id = str(uuid4())
-        msg_id, msg = create_incoming_message(cell_id)
-        op.process_incoming_message('shell', msg)
-        assert op.get_cell_id(msg_id) == cell_id
-        assert op.get_msg_id(cell_id) == msg_id
-        # # Run it again without clearing to ensure it self clears
-        cell_id = str(uuid4())
-        msg_id, msg = create_incoming_message(cell_id)
-        op.process_incoming_message('shell', msg)
-        assert op.get_cell_id(msg_id) == cell_id
-        assert op.get_msg_id(cell_id) == msg_id
+# TODO: Implement this
+def test_output_task():
+    pass

jupyter_server_documents/tests/test_outputs_manager.py

@@ -47,7 +47,8 @@ def test_stream():
         file_id = str(uuid4())
         cell_id = str(uuid4())
         for s in streams:
-            op.write_stream(file_id, cell_id, s)
+            placeholder = op.write_output(file_id, cell_id, s)
+            op.write_stream(file_id, cell_id, s, placeholder)
         assert op.get_stream(file_id, cell_id) == text
 
 def test_display_data():
@@ -91,3 +92,74 @@ def file_not_found():
             op.get_output('a','b',0)
         with pytest.raises(FileNotFoundError):
             op.get_stream('a','b')       
+
+
+def test__compute_output_index_basic():
+    """
+    Test basic output index allocation for a cell without display ID
+    """
+    op = OutputsManager()
+    
+    # First output for a cell should be 0
+    assert op._compute_output_index('cell1') == 0
+    assert op._compute_output_index('cell1') == 1
+    assert op._compute_output_index('cell1') == 2
+
+def test__compute_output_index_with_display_id():
+    """
+    Test output index allocation with display IDs
+    """
+    op = OutputsManager()
+    
+    # First output for a cell with display ID
+    assert op._compute_output_index('cell1', 'display1') == 0
+    
+    # Subsequent calls with same display ID should return the same index
+    assert op._compute_output_index('cell1', 'display1') == 0
+    
+    # Different display ID should get a new index
+    assert op._compute_output_index('cell1', 'display2') == 1
+
+
+def test__compute_output_index_multiple_cells():
+    """
+    Test output index allocation across multiple cells
+    """
+    op = OutputsManager()
+    
+    assert op._compute_output_index('cell1') == 0
+    assert op._compute_output_index('cell1') == 1
+    assert op._compute_output_index('cell2') == 0
+    assert op._compute_output_index('cell2') == 1
+
+def test_display_id_index_retrieval():
+    """
+    Test retrieving output index for a display ID
+    """
+    op = OutputsManager()
+    
+    op._compute_output_index('cell1', 'display1')
+    
+    assert op.get_output_index('display1') == 0
+    assert op.get_output_index('non_existent_display') is None
+
+def test_display_ids():
+    """
+    Test tracking of display IDs for a cell
+    """
+    op = OutputsManager()
+    
+    # Allocate multiple display IDs for a cell
+    op._compute_output_index('cell1', 'display1')
+    op._compute_output_index('cell1', 'display2')
+    
+    # Verify display IDs are tracked
+    assert 'cell1' in op._display_ids_by_cell_id
+    assert set(op._display_ids_by_cell_id['cell1']) == {'display1', 'display2'}
+    
+    # Clear cell indices
+    op.clear('file1', 'cell1')
+    
+    # Verify display IDs are cleared
+    assert 'display1' not in op._display_ids_by_cell_id
+    assert 'display2' not in op._display_ids_by_cell_id