Refactor OutputsManager to align with default Jupyter behavior (#163)

* Add safe msg_id handling in kernel_client.

* Contrain to Python 3.13 in dev environment.

* Create _is_stream_output utility function

* Update stream_limit to 500 to avoid triggering too early

* Use set literal

* Update get_outputs to remove stream_limit logic

The stream_limit logic is being moved in this PR to the writing of outputs, so get_outputs can just return all outputs.

* Create _append_to_stream_file utility method.

* Update process_loaded_notebook to handle exclude_outputs

* Modify private _process_loaded methods to handle exclude_outputs

* Update process_saving_notebooks to handle exclude_outputs

* Add comment about placeholder outputs wrt nbformat

* Fix write to better handle stream outputs and stream_limit at write time

* Remove call to clear in saving logic

* Refactor OutputsManager and add experimental OptimizedOutputsManager

This commit introduces a cleaner architecture for handling notebook outputs
and adds an experimental optimized version that supports excluding outputs
from saved notebook files.

Core changes to OutputsManager:
- Extract private utility functions (_create_output_url, _create_output_placeholder)
- Add comprehensive docstrings to all methods
- Simplify write() method by removing stream_limit logic
- Improve error handling in get_outputs() to return empty list instead of raising
- Consolidate output processing logic into _process_outputs_from_cell()
- Add helper methods: _upgrade_notebook_format(), _ensure_cell_id()
- Always write full outputs to notebook files on save (traditional Jupyter behavior)
- Remove stream-specific handling and StreamAPIHandler route

New OptimizedOutputsManager:
- Extends base OutputsManager with exclude_outputs metadata flag support
- When exclude_outputs=True: outputs stored only in runtime, not in saved files
- When exclude_outputs=False/unset: full outputs included in saved files (default)
- Implements stream_limit (500) for large stream outputs with link placeholders
- Provides _append_to_stream_file() for efficient stream handling
- Stream API handler for accessing accumulated stream outputs

Other improvements:
- Add __all__ to outputs/__init__.py for cleaner exports
- Expand test coverage with comprehensive test suite
- Rename private methods for clarity (_process_loaded_excluded_outputs, etc.)
- Update yroom_file_api to use process_saving_notebook correctly

The OptimizedOutputsManager is currently experimental and disabled by default.
StreamAPIHandler route is commented out until the feature is ready for production.

Author: ellisonbg

dev-environment.yml

@@ -2,7 +2,7 @@ name: serverdocs
 channels:
   - conda-forge
 dependencies:
-  - python
+  - python=3.13
   - nodejs=22
   - uv
   - jupyterlab

jupyter_server_documents/kernels/kernel_client.py

@@ -102,14 +102,14 @@ async def stop_listening(self):
     def handle_incoming_message(self, channel_name: str, msg: list[bytes]):
         """
         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
@@ -119,32 +119,35 @@ def handle_incoming_message(self, channel_name: str, msg: list[bytes]):
         # source channel.
         header = self.session.unpack(msg[0])
         msg_id = header["msg_id"]
-        msg_type = header.get("msg_type")                
+        msg_type = header.get("msg_type")
         metadata = self.session.unpack(msg[2])
         cell_id = metadata.get("cellId")
-        
+
         # 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:
                 yroom.set_cell_awareness_state(cell_id, "busy")
-        
+
         self.message_cache.add({
             "msg_id": msg_id,
             "channel": channel_name,
             "cell_id": cell_id
         })
         channel = getattr(self, f"{channel_name}_channel")
+        if channel.socket is None:
+            self.log.error(f"Channel {channel_name} socket is None! Cannot send message. Channel alive: {channel.is_alive()}")
+            raise AttributeError(f"Channel {channel_name} socket is None")
         channel.session.send_raw(channel.socket, msg)
 
     def send_kernel_info(self):
-        """Sends a kernel info message on the shell channel. Useful 
+        """Sends a kernel info message on the shell channel. Useful
         for determining if the kernel is busy or idle.
         """
         msg = self.session.msg("kernel_info_request")
@@ -240,10 +243,16 @@ 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')
+
+        # Safely get parent message ID and data
+        parent_header = dmsg.get("parent_header", {})
+        parent_msg_id = parent_header.get("msg_id")
+
+        # Get parent message data from cache (may be None if not found)
+        parent_msg_data = self.message_cache.get(parent_msg_id) if parent_msg_id else None
+
+        # Safely extract cell_id
+        cell_id = parent_msg_data.get('cell_id') if parent_msg_data else None
 
         # Handle different message types using pattern matching
         match dmsg["msg_type"]:

jupyter_server_documents/outputs/__init__.py

@@ -1,3 +1,9 @@
 from .handlers import outputs_handlers
 from .manager import OutputsManager
 from .output_processor import OutputProcessor
+
+__all__ = [
+    'OutputsManager',
+    'OutputProcessor',
+    'outputs_handlers'
+]

jupyter_server_documents/outputs/handlers.py

@@ -86,5 +86,7 @@ async def get(self, file_id=None, cell_id=None):
 
 outputs_handlers = [
     (rf"/api/outputs/{_file_id_regex}/{_cell_id_regex}(?:/{_output_index_regex}.output)?", OutputsAPIHandler),
-    (rf"/api/outputs/{_file_id_regex}/{_cell_id_regex}/stream", StreamAPIHandler),
+    # We have disabled this for now as OptimizedOutputsManager is experimental.
+    # Uncomment this to use OptimizedOutputsManager.
+    # (rf"/api/outputs/{_file_id_regex}/{_cell_id_regex}/stream", StreamAPIHandler),
 ]