Multiple enhancements

.gitignore

@@ -101,3 +101,4 @@ dmypy.json
 # OSX files
 .DS_Store
 .jupyter_ystore.db
+.vscode

README.md

@@ -34,6 +34,83 @@ that the server extension is enabled:
 jupyter server extension list
 ```
 
+## How does it works
+
+### Generic case
+
+Execution of a Python code snippet: `print("hello")`
+
+```mermaid
+sequenceDiagram
+    Frontend->>+Server: POST /api/kernels/<id>/execute
+    Server->>+ExecutionStack: Create asyncio.Task
+    ExecutionStack->>Kernel: Execute request msg
+    activate Kernel
+    ExecutionStack-->>Server: Task uid
+    Server-->>-Frontend: Returns task uid
+    loop Running
+        Kernel->>Shared Document: Add output
+        Shared Document->>Frontend: Document update
+    end
+    loop While status is 202
+        Frontend->>+Server: GET /api/kernels/<id>/requests/<uid>
+        Server->>ExecutionStack: Get task result
+        ExecutionStack-->>Server: null
+        Server-->>-Frontend: Request status 202
+    end
+    Kernel-->>ExecutionStack: Execution reply
+    deactivate Kernel
+    Frontend->>+Server: GET /api/kernels/<id>/requests/<uid>
+    Server->>ExecutionStack: Get task result
+    ExecutionStack-->>Server: Result
+    Server-->>-Frontend: Status 200 & result
+```
+
+### With input case
+
+
+Execution of a Python code snippet: `input("Age:")`
+
+```mermaid
+sequenceDiagram
+    Frontend->>+Server: POST /api/kernels/<id>/execute
+    Server->>+ExecutionStack: Create asyncio.Task
+    ExecutionStack->>Kernel: Execute request msg
+    activate Kernel
+    ExecutionStack-->>Server: Task uid
+    Server-->>-Frontend: Returns task uid
+    loop Running
+        Kernel->>Shared Document: Add output
+        Shared Document->>Frontend: Document update
+    end
+    loop While status is 202
+        Frontend->>+Server: GET /api/kernels/<id>/requests/<uid>
+        Server->>ExecutionStack: Get task result
+        ExecutionStack-->>Server: null
+        Server-->>-Frontend: Request status 202
+    end
+    Kernel->>ExecutionStack: Set pending input
+    Frontend->>+Server: GET /api/kernels/<id>/requests/<uid>
+    Server->>ExecutionStack: Get task result
+    ExecutionStack-->>Server: Pending input
+    Server-->>-Frontend: Status 300 & Pending input
+    Frontend->>+Server: POST /api/kernels/<id>/input
+    Server->>Kernel: Send input msg
+    Server-->>-Frontend: 
+    loop While status is 202
+        Frontend->>+Server: GET /api/kernels/<id>/requests/<uid>
+        Server->>ExecutionStack: Get task result
+        ExecutionStack-->>Server: null
+        Server-->>-Frontend: Request status 202
+    end
+    Kernel-->>ExecutionStack: Execution reply
+    deactivate Kernel
+    Frontend->>+Server: GET /api/kernels/<id>/requests/<uid>
+    Server->>ExecutionStack: Get task result
+    ExecutionStack-->>Server: Result
+    Server-->>-Frontend: Status 200 & result
+```
+
 ## Contributing
 
 ### Development install
@@ -46,7 +123,6 @@ jupyter server extension list
 pip install -e .
 ```
 
-
 You can watch the source directory and run your Jupyter Server-based application at the same time in different terminals to watch for changes in the extension's source and automatically rebuild the extension.  For example,
 when running JupyterLab:
 
@@ -61,6 +137,25 @@ server directly:
 jupyter server --autoreload
 ```
 
+### Manual testing
+
+```bash
+# Terminal 1.
+jupyter server --port 8888 --autoreload --ServerApp.disable_check_xsrf=True --IdentityProvider.token= --ServerApp.port_retries=0
+
+# Terminal 2.
+KERNEL=$(curl -X POST http://localhost:8888/api/kernels)
+echo $KERNEL
+KERNEL_ID=$(echo $KERNEL | jq --raw-output '.id')
+echo $KERNEL_ID
+REQUEST=$(curl --include http://localhost:8888/api/kernels/$KERNEL_ID/execute -d "{ \"code\": \"print('1+1')\" }")
+RESULT=$(echo $REQUEST | grep -i ^Location: | cut -d' ' -f2 | tr -d '\r')
+echo $RESULT
+
+curl http://localhost:8888$RESULT
+{"status": "ok", "execution_count": 1, "outputs": "[{\"output_type\": \"stream\", \"name\": \"stdout\", \"text\": \"1+1\\n\"}]"}
+```
+
 ### Running Tests
 
 Install dependencies:
@@ -78,7 +173,7 @@ pytest
 pytest jupyter_server_nbmodel/tests/test_handlers.py
 
 # To run a specific test
-pytest jupyter_server_nbmodel/tests/test_handlers.py -k "test_get"
+pytest jupyter_server_nbmodel/tests/test_handlers.py -k "test_post_execute"
 ```
 
 ### Development uninstall

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,51 @@
 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, InputHandler, 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}/input",
+                    InputHandler,
+                ),
+                (
+                    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