WIP: Add notebook observers

.gitignore

@@ -128,3 +128,5 @@ dmypy.json
 
 # For local testing
 playground/
+
+.vscode/

jupyter_ai_router/extension.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
+import asyncio
+import json
 from typing import TYPE_CHECKING
 import time
 from jupyter_events import EventLogger
 from jupyter_server.extension.application import ExtensionApp
+from jupyter_ydoc.ybasedoc import YBaseDoc
 
 from jupyter_ai_router.handlers import RouteHandler
 
@@ -41,6 +44,18 @@ class RouterExtension(ExtensionApp):
 
     router: MessageRouter
 
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        """
+        Returns a reference to the asyncio event loop.
+        """
+        return asyncio.get_event_loop_policy().get_event_loop()
+
+    @property
+    def fileid_manager(self):
+        return self.serverapp.web_app.settings["file_id_manager"]
+
+
     def initialize_settings(self):
         """Initialize router settings and event listeners."""
         start = time.time()
@@ -59,10 +74,32 @@ def initialize_settings(self):
             self.event_logger.add_listener(
                 schema_id=JUPYTER_COLLABORATION_EVENTS_URI, listener=self._on_chat_event
             )
-
+            
         elapsed = time.time() - start
         self.log.info(f"Initialized RouterExtension in {elapsed:.2f}s")
 
+    def _get_global_awareness(self):
+        # TODO: make this compatible with jcollab
+        jcollab_api = self.serverapp.web_app.settings["jupyter_server_ydoc"]
+        yroom_manager = jcollab_api.yroom_manager
+        yroom = yroom_manager.get_room("JupyterLab:globalAwareness")
+        return yroom.get_awareness()
+
+    async def _room_id_from_path(self, path: str) -> str | None:
+        """Returns room_id from document path"""
+        # TODO: Make this compatible with jcollab
+        yroom_manager = self.serverapp.web_app.settings["yroom_manager"]
+        for room_id in yroom_manager._rooms_by_id:
+            if room_id == "JupyterLab:globalAwareness":
+                continue
+            ydoc = await self._get_doc(room_id)
+            state = ydoc.awareness.get_local_state()
+            file_id = state["file_id"]
+            ydoc_path = self.fileid_manager.get_path(file_id)
+            if ydoc_path == path:
+                print(f"Found match in path {path}")
+                return room_id
+
     async def _on_chat_event(
         self, logger: EventLogger, schema_id: str, data: dict
     ) -> None:
@@ -87,6 +124,84 @@ async def _on_chat_event(
         # Connect chat to router
         self.router.connect_chat(room_id, ychat)
 
+    async def _on_notebook_event(
+        self, logger: EventLogger, schema_id: str, data: dict
+    ) -> None:
+        """Handle notebook room events and connect new chats to router."""
+        # Only handle notebook room initialization events
+        if not (
+            data["room"].startswith("json:notebook:")
+            and data["action"] == "initialize"
+            and data["msg"] == "Room initialized"
+        ):
+            return
+
+        room_id = data["room"]
+        self.log.info(f"New notebook room detected: {room_id}")
+
+        # Get YDoc document for the room
+        ydoc = await self._get_doc(room_id)
+        if ydoc is None:
+            self.log.error(f"Failed to get YDoc for room {room_id}")
+            return
+
+        # Connect notebook to router
+        self.router.connect_notebook(room_id, ydoc)
+
+    async def _get_doc(self, room_id: str) -> YBaseDoc | None:
+        """
+        Get YDoc instance for a room ID.
+
+        Dispatches to either `_get_doc_jcollab()` or `_get_doc_jsd()` based on
+        whether `jupyter_server_documents` is installed.
+        """
+
+        if JSD_PRESENT:
+            return await self._get_doc_jsd(room_id)
+        else:
+            return await self._get_doc_jcollab(room_id)
+
+    async def _get_doc_jcollab(self, room_id: str) -> YBaseDoc | None:
+        """
+        Method used to retrieve the `YDoc` instance for a given room when
+        `jupyter_server_documents` **is not** installed.
+        """
+        if not self.serverapp:
+            return None
+
+        try:
+            collaboration = self.serverapp.web_app.settings["jupyter_server_ydoc"]
+            document = await collaboration.get_document(room_id=room_id, copy=False)
+            return document
+        except Exception as e:
+            self.log.error(f"Error getting ydoc for {room_id}: {e}")
+            return None
+
+    async def _get_doc_jsd(self, room_id: str) -> YBaseDoc | None:
+        """
+        Method used to retrieve the `YDoc` instance for a given room when
+        `jupyter_server_documents` **is** installed.
+
+        This method uniquely attaches a callback which is fired whenever the
+        `YDoc` is reset.
+        """
+        if not self.serverapp:
+            return None
+
+        try:
+            jcollab_api = self.serverapp.web_app.settings["jupyter_server_ydoc"]
+            yroom_manager = jcollab_api.yroom_manager
+            yroom = yroom_manager.get_room(room_id)
+
+            def _on_ydoc_reset(new_ydoc: YBaseDoc):
+                self.router._on_notebook_reset(room_id, new_ydoc)
+
+            ydoc = await yroom.get_jupyter_ydoc(on_reset=_on_ydoc_reset)
+            return ydoc
+        except Exception as e:
+            self.log.error(f"Error getting ydoc for {room_id}: {e}")
+            return None
+
     async def _get_chat(self, room_id: str) -> YChat | None:
         """
         Get YChat instance for a room ID.