Use execution stack

Author: fcollonval

conftest.py

@@ -2,6 +2,11 @@
 
 pytest_plugins = ["jupyter_server.pytest_plugin"]
 
+
 @pytest.fixture
 def jp_server_config(jp_server_config):
-     return {"ServerApp": {"jpserver_extensions": {"jupyter_server_nbmodel": True, "jupyter_server_ydoc": True}}}
+    return {
+        "ServerApp": {
+            "jpserver_extensions": {"jupyter_server_nbmodel": True, "jupyter_server_ydoc": True}
+        }
+    }

jupyter_server_config.py

@@ -1 +1 @@
-c.YDocExtension.server_side_execution = True
+c.YDocExtension.server_side_execution = True  # noqa F821

jupyter_server_nbmodel/__init__.py

@@ -1,10 +1,9 @@
 """A Jupyter Server extension to execute code cell from the server."""
+
 from .extension import Extension
+
 __version__ = "0.1.0"
 
 
 def _jupyter_server_extension_points():
-    return [{
-        "module": "jupyter_server_nbmodel",
-        "app": Extension
-    }]
+    return [{"module": "jupyter_server_nbmodel", "app": Extension}]

jupyter_server_nbmodel/extension.py

@@ -3,23 +3,47 @@
 from jupyter_server.extension.application import ExtensionApp
 from jupyter_server.services.kernels.handlers import _kernel_id_regex
 
-from .handlers import ExecuteHandler
+from .handlers import ExecuteHandler, ExecutionStack, RequestHandler
 from .log import get_logger
 
 RTC_EXTENSIONAPP_NAME = "jupyter_server_ydoc"
 
+STOP_TIMEOUT = 3
+
+_request_id_regex = r"(?P<request_id>\w+-\w+-\w+-\w+-\w+)"
+
 
 class Extension(ExtensionApp):
     name = "jupyter_server_nbmodel"
 
     def initialize_handlers(self):
         rtc_extension = None
-        rtc_extensions = self.serverapp.extension_manager.extension_apps.get(RTC_EXTENSIONAPP_NAME, set())
+        rtc_extensions = self.serverapp.extension_manager.extension_apps.get(
+            RTC_EXTENSIONAPP_NAME, set()
+        )
         n_extensions = len(rtc_extensions)
         if n_extensions:
             if n_extensions > 1:
                 get_logger().warning("%i collaboration extensions found.", n_extensions)
             rtc_extension = next(iter(rtc_extensions))
-        self.handlers.extend([
-            (f"/api/kernels/{_kernel_id_regex}/execute", ExecuteHandler, { "ydoc_extension": rtc_extension })
-        ])
+
+        self.__tasks = ExecutionStack()
+
+        self.handlers.extend(
+            [
+                (
+                    f"/api/kernels/{_kernel_id_regex}/execute",
+                    ExecuteHandler,
+                    {"ydoc_extension": rtc_extension, "execution_stack": self.__tasks},
+                ),
+                (
+                    f"/api/kernels/{_kernel_id_regex}/requests/{_request_id_regex}",
+                    RequestHandler,
+                    {"execution_stack": self.__tasks},
+                ),
+            ]
+        )
+
+    async def stop_extension(self):
+        if hasattr(self, "__tasks"):
+            del self.__tasks