Move cell execution state to the document's awareness object

Author: Zsailer

jupyter_server_documents/kernels/kernel_client.py

@@ -100,21 +100,45 @@ async def stop_listening(self):
     _listening_task: t.Optional[t.Awaitable] = Any(allow_none=True)
 
     def handle_incoming_message(self, channel_name: str, msg: list[bytes]):
-        """Use the given session to send the message."""
+        """
+        Handle incoming kernel messages and set up immediate cell execution state tracking.
+        
+        This method processes incoming kernel messages and caches them for response mapping.
+        Importantly, it detects execute_request messages and immediately sets the corresponding
+        cell state to 'busy' to provide real-time feedback for queued cell executions.
+        
+        This ensures that when multiple cells are executed simultaneously, all queued cells
+        show a '*' prompt immediately, not just the currently executing cell.
+        
+        Args:
+            channel_name: The kernel channel name (shell, iopub, etc.)
+            msg: The raw kernel message as bytes
+        """
         # Cache the message ID and its socket name so that
         # any response message can be mapped back to the
         # source channel.
         header = self.session.unpack(msg[0])
-        msg_id = header["msg_id"]                
+        msg_id = header["msg_id"]
+        msg_type = header.get("msg_type")                
         metadata = self.session.unpack(msg[2])
         cell_id = metadata.get("cellId")
 
-        # Clear cell outputs if cell is re-executedq
+        # Clear cell outputs if cell is re-executed
         if cell_id:
             existing = self.message_cache.get(cell_id=cell_id)
             if existing and existing['msg_id'] != msg_id:
                 asyncio.create_task(self.output_processor.clear_cell_outputs(cell_id))
 
+        # IMPORTANT: Set cell to 'busy' immediately when execute_request is received
+        # This ensures queued cells show '*' prompt even before kernel starts processing them
+        if msg_type == "execute_request" and channel_name == "shell" and cell_id:
+            for yroom in self._yrooms:
+                awareness = yroom.get_awareness()
+                if awareness is not None:
+                    cell_states = awareness.get_local_state().get("cell_execution_states", {})
+                    cell_states[cell_id] = "busy"
+                    awareness.set_local_state_field("cell_execution_states", cell_states)
+        
         self.message_cache.add({
             "msg_id": msg_id,
             "channel": channel_name,
@@ -240,27 +264,27 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
                     metadata["metadata"]["language_info"] = language_info
 
             case "status":
-                # Unpack cell-specific information and determine execution state
+                # Handle kernel status messages and update cell execution states
+                # This provides real-time feedback about cell execution progress
                 content = self.session.unpack(dmsg["content"])
                 execution_state = content.get("execution_state")
+                
                 # Update status across all collaborative rooms
                 for yroom in self._yrooms:
-                    # If this status came from the shell channel, update
-                    # the notebook status.
-                    if parent_msg_data["channel"] == "shell":                     
-                        awareness = yroom.get_awareness()
-                        if awareness is not None:
+                    awareness = yroom.get_awareness()
+                    if awareness is not None:
+                        # If this status came from the shell channel, update
+                        # the notebook kernel status.
+                        if parent_msg_data and parent_msg_data.get("channel") == "shell":                     
                             # Update the kernel execution state at the top document level
                             awareness.set_local_state_field("kernel", {"execution_state": execution_state})
-                    # Specifically update the running cell's execution state if cell_id is provided
-                    if cell_id:
-                        notebook = await yroom.get_jupyter_ydoc()
-                        _, target_cell = notebook.find_cell(cell_id)
-                        if target_cell:
-                            # Adjust state naming convention from 'busy' to 'running' as per JupyterLab expectation
-                            # https://github.com/jupyterlab/jupyterlab/blob/0ad84d93be9cb1318d749ffda27fbcd013304d50/packages/cells/src/widget.ts#L1670-L1678
-                            state = 'running' if execution_state == 'busy' else execution_state
-                            target_cell["execution_state"] = state
+                        
+                        # Store cell execution state for persistence across client connections
+                        # This ensures that cell execution states survive page refreshes
+                        if cell_id:
+                            cell_states = awareness.get_local_state().get("cell_execution_states", {})
+                            cell_states[cell_id] = execution_state
+                            awareness.set_local_state_field("cell_execution_states", cell_states)
                             break
 
             case "execute_input":
@@ -278,8 +302,7 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
             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')
+                    if self.output_processor:
                         content = self.session.unpack(dmsg["content"])
                         self.output_processor.process_output(dmsg['msg_type'], cell_id, content)
 

jupyter_server_documents/rooms/yroom.py

@@ -369,6 +369,23 @@ def get_awareness(self) -> pycrdt.Awareness:
         """
         return self._awareness
 
+    def get_cell_execution_states(self) -> dict:
+        """
+        Returns the persistent cell execution states for this room.
+        These states survive client disconnections but are not saved to disk.
+        """
+        if not hasattr(self, '_cell_execution_states'):
+            self._cell_execution_states = {}
+        return self._cell_execution_states
+    
+    def set_cell_execution_state(self, cell_id: str, execution_state: str) -> None:
+        """
+        Sets the execution state for a specific cell.
+        This state persists across client disconnections.
+        """
+        if not hasattr(self, '_cell_execution_states'):
+            self._cell_execution_states = {}
+        self._cell_execution_states[cell_id] = execution_state
 
     def add_message(self, client_id: str, message: bytes) -> None:
         """
@@ -512,7 +529,7 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
             return
 
         self.clients.mark_synced(client_id)
-        
+
         # Send SyncStep1 message
         try:
             assert isinstance(new_client.websocket, WebSocketHandler)

jupyter_server_documents/tests/kernels/test_kernel_client_integration.py

@@ -20,6 +20,17 @@ def mock_yroom_with_notebook(self):
         # Create a real YDoc and YNotebook
         ydoc = pycrdt.Doc()
         awareness = MagicMock(spec=pycrdt.Awareness)  # Mock awareness instead of using real one
+        
+        # Mock the local state to track changes
+        local_state = {}
+        awareness.get_local_state = MagicMock(return_value=local_state)
+        
+        # Mock set_local_state_field to actually update the local_state dict
+        def mock_set_local_state_field(field, value):
+            local_state[field] = value
+        
+        awareness.set_local_state_field = MagicMock(side_effect=mock_set_local_state_field)
+        
         ynotebook = YNotebook(ydoc, awareness)
 
         # Add a simple notebook structure with one cell
@@ -54,6 +65,18 @@ def mock_yroom_with_notebook(self):
         yroom.get_jupyter_ydoc = AsyncMock(return_value=ynotebook)
         yroom.get_awareness = MagicMock(return_value=awareness)
 
+        # Add persistent cell execution state storage
+        yroom._cell_execution_states = {}
+        
+        def mock_get_cell_execution_states():
+            return yroom._cell_execution_states
+        
+        def mock_set_cell_execution_state(cell_id, execution_state):
+            yroom._cell_execution_states[cell_id] = execution_state
+        
+        yroom.get_cell_execution_states = MagicMock(side_effect=mock_get_cell_execution_states)
+        yroom.set_cell_execution_state = MagicMock(side_effect=mock_set_cell_execution_state)
+        
         return yroom, ynotebook
 
     @pytest.fixture
@@ -109,7 +132,7 @@ async def test_execute_input_updates_execution_count(self, kernel_client_with_yr
 
     @pytest.mark.asyncio
     async def test_status_message_updates_cell_execution_state(self, kernel_client_with_yroom):
-        """Test that status messages update cell execution state in YDoc."""
+        """Test that status messages update cell execution state in YRoom for persistence and awareness for real-time updates."""
         client, yroom, ynotebook = kernel_client_with_yroom
 
         # Mock message cache to return cell_id and channel
@@ -129,11 +152,16 @@ async def test_status_message_updates_cell_execution_state(self, kernel_client_w
         # Process the message
         await client.handle_document_related_message(msg_parts[1:])  # Skip delimiter
 
-        # Verify the cell execution state was updated to 'running' (converted from 'busy')
-        cells = ynotebook.get_cell_list()
-        target_cell = next((cell for cell in cells if cell.get("id") == cell_id), None)
-        assert target_cell is not None
-        assert target_cell.get("execution_state") == "running"
+        # Verify the cell execution state was stored in YRoom for persistence
+        cell_states = yroom.get_cell_execution_states()
+        assert cell_states[cell_id] == "busy"
+        
+        # Verify it's also in awareness for real-time updates
+        awareness = yroom.get_awareness()
+        local_state = awareness.get_local_state()
+        assert local_state is not None
+        assert "cell_execution_states" in local_state
+        assert local_state["cell_execution_states"][cell_id] == "busy"
 
     @pytest.mark.asyncio
     async def test_kernel_info_reply_updates_language_info(self, kernel_client_with_yroom):
@@ -304,12 +332,16 @@ async def test_complete_execution_flow(self, kernel_client_with_yroom):
         )
         await client.handle_document_related_message(msg_parts[1:])
 
-        # Verify final state of the cell in YDoc
+        # Verify final state of the cell in YDoc and awareness
         cells = ynotebook.get_cell_list()
         target_cell = next((cell for cell in cells if cell.get("id") == cell_id), None)
         assert target_cell is not None
         assert target_cell.get("execution_count") == 1
-        assert target_cell.get("execution_state") == "idle"
+        
+        # Verify execution state is stored in awareness, not YDoc
+        awareness = yroom.get_awareness()
+        cell_execution_states = awareness.get_local_state().get("cell_execution_states", {})
+        assert cell_execution_states.get(cell_id) == "idle"
 
         # Verify output processor was called for the result
         client.output_processor.process_output.assert_called_with(
@@ -383,15 +415,12 @@ def mock_get(msg_id):
         )
         await client.handle_document_related_message(msg_parts2[1:])
 
-        # Verify both cells have correct states
-        cells = ynotebook.get_cell_list()
-        cell1 = next((cell for cell in cells if cell.get("id") == "test-cell-1"), None)
-        cell2 = next((cell for cell in cells if cell.get("id") == "test-cell-2"), None)
+        # Verify both cells have correct states in awareness
+        awareness = yroom.get_awareness()
+        cell_execution_states = awareness.get_local_state().get("cell_execution_states", {})
 
-        assert cell1 is not None
-        assert cell1.get("execution_state") == "running"  # 'busy' -> 'running'
-        assert cell2 is not None
-        assert cell2.get("execution_state") == "idle"
+        assert cell_execution_states.get("test-cell-1") == "busy"  # 'busy' state
+        assert cell_execution_states.get("test-cell-2") == "idle"
 
     @pytest.mark.asyncio
     async def test_message_without_cell_id_skips_cell_updates(self, kernel_client_with_yroom):