Work to integrate output processor into kernel client.

Author: ellisonbg

jupyter_rtc_core/kernels/kernel_client.py

@@ -1,16 +1,17 @@
 """
 A new Kernel client that is aware of ydocuments.
 """
+import anyio
 import asyncio
 import json
 import typing as t
-from traitlets import Set
-from traitlets import Instance
-from traitlets import Any
-from .utils import LRUCache
+
+from traitlets import Set, Instance, Any, Type, default
 from jupyter_client.asynchronous.client import AsyncKernelClient
-import anyio
+
+from .utils import LRUCache
 from jupyter_rtc_core.rooms.yroom import YRoom
+from jupyter_rtc_core.outputs import OutputProcessor
 
 
 class DocumentAwareKernelClient(AsyncKernelClient): 
@@ -35,6 +36,20 @@ class DocumentAwareKernelClient(AsyncKernelClient):
     # status messages.
     _yrooms: t.Set[YRoom] = Set(trait=Instance(YRoom), default_value=set())
 
+    output_processor = Instance(
+        OutputProcessor,
+        allow_none=True
+    )
+
+    output_process_class = Type(
+        klass=OutputProcessor,
+        default_value=OutputProcessor
+    ).tag(config=True)
+
+    @default("output_processor")
+    def _default_output_processor(self) -> OutputProcessor:
+        self.log.info("Creating output processor")
+        return OutputProcessor(parent=self, config=self.config)
 
     async def start_listening(self):
         """Start listening to messages coming from the kernel.
@@ -78,7 +93,8 @@ def handle_incoming_message(self, channel_name: str, msg: list[bytes]):
         # Cache the message ID and its socket name so that
         # any response message can be mapped back to the
         # source channel.
-        header = header = json.loads(msg[0])
+        self.output_processor.process_incoming_message(channel=channel_name, msg=msg)
+        header = header = json.loads(msg[0]) # TODO: use session.unpack
         msg_id = header["msg_id"]
         self.message_source_cache[msg_id] = channel_name
         channel = getattr(self, f"{channel_name}_channel")
@@ -159,10 +175,9 @@ async def handle_outgoing_message(self, channel_name: str, msg: list[bytes]):
         # Intercept messages that are IOPub focused.
         if channel_name == "iopub":
             message_returned = await self.handle_iopub_message(msg)
-            # TODO: If the message is not returned by the iopub handler, then
+            # If the message is not returned by the iopub handler, then
             # return here and do not forward to listeners.
             if not message_returned:
-                self.log.warn(f"If message is handled donot forward after adding output manager")
                 return
 
         # Update the last activity.
@@ -182,47 +197,38 @@ async def handle_iopub_message(self, msg: list[bytes]) -> t.Optional[list[bytes]
         Returns
         ------- 
         Returns the message if it should be forwarded to listeners. Otherwise,
-        returns `None` and keeps (i.e. intercepts) the message from going
-        to listenres.
+        returns `None` and prevents (i.e. intercepts) the message from going
+        to listeners.
         """
-        # NOTE: Here's where we will inject the kernel state
-        # into the awareness of a document.
 
         try:
             dmsg = self.session.deserialize(msg, content=False)
         except Exception as e:
             self.log.error(f"Error deserializing message: {e}")
             raise ValueError
 
-        if dmsg["msg_type"] == "status":
-            # Forward to all yrooms.
-            for yroom in self._yrooms:
-                # NOTE: We need to create a real message here.
-                awareness_update_message = b""
-                self.log.debug(f"Update Awareness here: {dmsg}. YRoom: {yroom}")
-                #self.log.debug(f"Getting YDoc: {await yroom.get_ydoc()}")
-                #yroom.add_message(awareness_update_message)
+        # if dmsg["msg_type"] == "status":
+        #     # Forward to all yrooms.
+        #     for yroom in self._yrooms:
+        #         # NOTE: We need to create a real message here.
+        #         awareness_update_message = b""
+        #         self.log.debug(f"Update Awareness here: {dmsg}. YRoom: {yroom}")
+        #         #self.log.debug(f"Getting YDoc: {await yroom.get_ydoc()}")
+        #         #yroom.add_message(awareness_update_message)
 
-            # TODO: returning message temporarily to not break UI
-            return msg
+        #     # TODO: returning message temporarily to not break UI
+        #     return msg
 
-
-        # NOTE: Inject display data into ydoc.
-        if dmsg["msg_type"] == "display_data":
-            # Forward to all yrooms.
-            for yroom in self._yrooms:
-                update_document_message = b""
-                self.log.debug(f"Update Document here: {dmsg}. Yroom: {yroom}")
-                #self.log.debug(f"Getting YDoc: {await yroom.get_ydoc()}")
-                #yroom.add_message(update_document_message)
-            
-            # TODO: returning message temporarily to not break UI
+        if self.output_processor is not None and dmsg["msg_type"] in ("stream", "display_data", "execute_result", "error"):
+            dmsg = self.output_processor.process_outgoing_message(dmsg)
+
+        # If process_outgoing_message returns None, return None so the message isn't
+        # sent to clients, otherwise return the original serialized message.
+        if dmsg is None:
+            return None
+        else:
             return msg
 
-        # If the message isn't handled above, return it and it will
-        # be forwarded to all listeners
-        return msg
-
     async def add_yroom(self, yroom: YRoom):
         """
         Register a YRoom with this kernel client. YRooms will

jupyter_rtc_core/outputs/manager.py

@@ -99,7 +99,7 @@ def write_stream(self, file_id, cell_id, output) -> Map:
         # Now create the placeholder output
         if count < self.stream_limit:
             # Return the original if we haven't reached the limit
-            placeholder = Map(output)
+            placeholder = Map({})
         elif count == self.stream_limit:
             # Return a link to the full stream output
             placeholder = Map({

jupyter_rtc_core/outputs/output_processor.py

@@ -15,9 +15,8 @@ class OutputProcessor(LoggingConfigurable):
 
     use_outputs_service = Bool(
         default_value=True,
-        config=True,
         help="Should outputs be routed to the outputs service to minimize the in memory ydoc size."
-    )
+    ).tag(config=True)
 
     @property
     def settings(self):
@@ -107,35 +106,58 @@ def process_incoming_message(self, channel: str, msg: list[bytes]):
             if self._file_id is not None:
                 if self.use_outputs_service:
                     self.outputs_manager.clear(file_id=self._file_id, cell_id=cell_id)
+                    self.outputs_manager.clear(file_id=self._file_id)
         self.log.info(f"Saving (msg_id, cell_id): ({msg_id} {cell_id})")
         self.set_cell_id(msg_id, cell_id)
 
     # Outgoing messages
 
-    def process_outgoing_message(self, channel: str, msg: list[bytes]):
-        """Process outgoing messagers from the kernel."""
-        dmsg = self.kernel_client.session.deserialize(msg)
+    def process_outgoing_message(self, dmsg: dict):
+        """Process outgoing messages from the kernel.
+        
+        This returns the input dmsg if no the message should be sent to
+        clients, or None if it should not be sent.
+
+        The dmsg is a deserialized message generated by calling:
+
+        > self.kernel_client.session.deserialize(dmsg, content=False)
+
+        The content has not been deserialized yet as we need to verify we
+        should process it.
+        """
+        self.log.info("Process outgoing message")
         msg_type = dmsg["header"]["msg_type"]
+        self.log.info(f"msg_type: {msg_type}")
+        if msg_type not in ("stream", "display_data", "execute_result", "error"):
+            return dmsg
         msg_id = dmsg["parent_header"]["msg_id"]
-        content = dmsg["content"]
+        self.log.info(f"msg_type: {msg_type}")
+        content = self.parent.session.unpack(dmsg["content"])
+        self.log.info(f"content: {content}")
         cell_id = self.get_cell_id(msg_id)
+        self.log.info(f"cell_id: {cell_id}")
         if cell_id is None:
             # This is valid as cell_id is optional
-            return
+            return dmsg
         asyncio.create_task(self.output_task(msg_type, cell_id, content))
         return None # Don't allow the original message to propagate to the frontend
 
     async def output_task(self, msg_type, cell_id, content):
         """A coroutine to handle output messages."""
+        self.log.info(f"output_task: {msg_type} {cell_id}")
         try:
             kernel_session = await self.session_manager.get_session(kernel_id=self.kernel_id)
-        except: # TODO: what exception to catch and log?
+        except Exception as e:
+            self.log.error(f"An exception occurred when processing output for cell {cell_id}")
+            self.log.exception(e)
             return
         else:
             path = kernel_session["path"]
 
         file_id = self.file_id_manager.get_id(path)
+        self.log.info(f"Output for file_id, cell_id: {file_id} {cell_id}")
         if file_id is None:
+            self.log.error(f"Could not find file_id for path: {path}")
             return
         self._file_id = file_id
         try:
@@ -145,7 +167,9 @@ async def output_task(self, msg_type, cell_id, content):
                 file_format='json',
                 content_type='notebook'
             )
-        except: # TODO: what exception to catch and log?
+        except Exception as e:
+            self.log.error(f"An exception occurred when processing output for cell {cell_id}")
+            self.log.exception(e)
             return
         cells = notebook.ycells