Repository: datalayer/jupyter-server-nbmodel

Issue:
Should code be executed asynchronously in RTC mode?
<!--
Welcome! Thanks for thinking of a way to improve Datalayer. If this solves a problem for you, then it probably solves that problem for lots of people! So the whole community will benefit from this request.

Before creating a new feature request please search the issues for relevant feature requests.
-->

### Problem

<!-- Provide a clear and concise description of what problem this feature will solve. For example:

* I'm always frustrated when [...] because [...]
* I would like it if [...] happened when I [...] because [...]
-->

https://github.com/datalayer/jupyter-server-nbmodel/blob/b122d68e2c9608f24201aa757a47ea922147825d/jupyter_server_nbmodel/handlers.py#L125-L131

I've found that current implementations where cell takes longer to execute can lead to http request timeouts (especially behind some reverse proxies), and it's difficult for the front-end to determine if the request to execute the code was successfully sent.

### Proposed Solution

<!-- Provide a clear and concise description of a way to accomplish what you want. For example:

* Add an option so that when [...]  [...] will happen
 -->

1. We can simply support asynchronous mode: do execute and safe in an asyncio.Task.
2. Maybe we need an optimistic lock to prevent racing  when concurrent write to the ynotebook.

### Additional context

<!-- Add any other context or screenshots about the feature request here. You can also include links to examples of other programs that have something similar to your request. For example:

* Another project [...] solved this by [...]
-->


Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 1048b01..795efd7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -101,3 +101,4 @@ dmypy.json
 # OSX files
 .DS_Store
 .jupyter_ystore.db
+.vscode
diff --git a/README.md b/README.md
index 7916d4a..4934ec9 100644
--- a/README.md
+++ b/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
diff --git a/jupyter_server_config.py b/jupyter_server_config.py
index e47f652..5df935b 100644
--- a/jupyter_server_config.py
+++ b/jupyter_server_config.py
@@ -1,1 +1,1 @@
-c.YDocExtension.server_side_execution = True
+c.YDocExtension.server_side_execution = True  # noqa F821
diff --git a/jupyter_server_nbmodel/__init__.py b/jupyter_server_nbmodel/__init__.py
index 19e5edb..bc0bb55 100644
--- a/jupyter_server_nbmodel/__init__.py
+++ b/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}]
diff --git a/jupyter_server_nbmodel/extension.py b/jupyter_server_nbmodel/extension.py
index ac33ba0..2a5b273 100644
--- a/jupyter_server_nbmodel/extension.py
+++ b/jupyter_server_nbmodel/extension.py
@@ -3,23 +3,51 @@ from __future__ import annotations
 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
diff --git a/jupyter_server_nbmodel/handlers.py b/jupyter_server_nbmodel/handlers.py
index 9003152..84db4d7 100644
--- a/jupyter_server_nbmodel/handlers.py
+++ b/jupyter_server_nbmodel/handlers.py
@@ -1,7 +1,13 @@
 from __future__ import annotations
 
+import asyncio
 import json
+import os
+import sys
+import traceback
 import typing as t
+import uuid
+from functools import partial
 from http import HTTPStatus
 
 import nbformat
@@ -13,27 +19,240 @@ from jupyter_server.extension.handler import ExtensionHandlerMixin
 from .log import get_logger
 
 if t.TYPE_CHECKING:
+    import jupyter_client
+
     try:
-        import pycrdt as y
         import jupyter_server_ydoc
+        import pycrdt as y
         from jupyter_ydoc.ynotebook import YNotebook
     except ImportError:
         # optional dependencies
         ...
 
 
+class ExecutionStack:
+    """Execution request stack.
+
+    It is keeping track of the execution requests.
+
+    The request result can only be queried once.
+    """
+
+    def __init__(self):
+        self.__pending_inputs: dict[str, dict] = {}
+        self.__tasks: dict[str, asyncio.Task] = {}
+
+    def __del__(self):
+        for task in filter(lambda t: not t.cancelled(), self.__tasks.values()):
+            task.cancel()
+
+    # FIXME
+    # def cancel(self, kernel_id: str, uid: str) -> None:
+    #     """Cancel the request ``uid``.
+
+    #     Args:
+    #         kernel_id : Kernel identifier
+    #         uid : Request identifier
+    #     """
+    #     get_logger().debug(f"Cancel request {uid}.")
+    #     if uid not in self.__tasks:
+    #         raise ValueError(f"Request {uid} does not exists.")
+
+    #     self.__tasks[uid].cancel()
+
+    def get(self, kernel_id: str, uid: str) -> t.Any:
+        """Get the request ``uid`` results, its pending input or None.
+
+        Args:
+            kernel_id : Kernel identifier
+            uid : Request identifier
+
+        Returns:
+            Any: None if the request is pending else its result or the pending input.
+
+        Raises:
+            ValueError: If the request ``uid`` does not exists.
+            asyncio.CancelledError: If the request ``uid`` was cancelled.
+        """
+        if uid not in self.__tasks:
+            raise ValueError(f"Request {uid} does not exists.")
+
+        if kernel_id in self.__pending_inputs:
+            return self.__pending_inputs.pop(kernel_id)
+
+        if self.__tasks[uid].done():
+            task = self.__tasks.pop(uid)
+            return task.result()
+        else:
+            return None
+
+    def put(
+        self, km: jupyter_client.manager.KernelManager, snippet: str, ycell: y.Map | None
+    ) -> str:
+        """Add a asynchronous execution request.
+
+        Args:
+            km: Kernel manager
+            snippet: Snippet to be executed
+            ycell: [optional] The cell to update with the execution results
+
+        Returns:
+            Request identifier
+        """
+        uid = str(uuid.uuid4())
+
+        self.__tasks[uid] = asyncio.create_task(
+            _execute_task(uid, km, snippet, ycell, partial(self._stdin_hook, km.kernel_id))
+        )
+        return uid
+
+    def _stdin_hook(self, kernel_id: str, msg: dict) -> None:
+        """Callback on stdin message.
+
+        It will register the pending input as temporary answer to the execution request.
+        """
+        get_logger().debug(f"Execution request {kernel_id} received a input request.")
+        if kernel_id in self.__pending_inputs:
+            get_logger().error(
+                f"Execution request {kernel_id} received a input request while waiting for an input.\n{msg}"  # noqa: E501
+            )
+
+        header = msg["header"].copy()
+        header["date"] = header["date"].isoformat()
+        self.__pending_inputs[kernel_id] = {
+            "parent_header": header,
+            "input_request": msg["content"],
+        }
+
+
+async def _execute_task(
+    uid,
+    km: jupyter_client.manager.KernelManager,
+    snippet: str,
+    ycell: y.Map | None,
+    stdin_hook: t.Callable[[dict], None] | None,
+) -> t.Any:
+    try:
+        get_logger().debug(f"Will execute request {uid}.")
+        result = await _execute_snippet(uid, km, snippet, ycell, stdin_hook)
+    except asyncio.CancelledError:
+        raise
+    except Exception as e:
+        exception_type, _, tb = sys.exc_info()
+        result = {
+            "type": exception_type.__qualname__,
+            "error": str(e),
+            "message": repr(e),
+            "traceback": traceback.format_tb(tb),
+        }
+        get_logger().error("Error for request %s.", result)
+    else:
+        get_logger().debug(f"Has executed request {uid}.")
+
+    return result
+
+
+async def _execute_snippet(
+    uid: str,
+    km: jupyter_client.manager.KernelManager,
+    snippet: str,
+    ycell: y.Map | None,
+    stdin_hook: t.Callable[[dict], None] | None,
+) -> dict[str, t.Any]:
+
+    if ycell is not None:
+        # Reset cell
+        with ycell.doc.transaction():
+            del ycell["outputs"][:]
+            ycell["execution_count"] = None
+
+    outputs = []
+    client = km.client()
+    client.session.session = uid
+    # FIXME
+    # client.session.username = username
+
+    # FIXME we don't check if the session is consistent (aka the kernel is linked to the document)
+    #   - should we?
+    try:
+        reply = await ensure_async(
+            client.execute_interactive(
+                snippet,
+                output_hook=partial(_output_hook, ycell, outputs),
+                stdin_hook=stdin_hook if client.allow_stdin else None,
+            )
+        )
+
+        reply_content = reply["content"]
+
+        if ycell is not None:
+            ycell["execution_count"] = reply_content.get("execution_count")
+
+        return {
+            "status": reply_content["status"],
+            "execution_count": reply_content.get("execution_count"),
+            # FIXME quid for buffers
+            "outputs": json.dumps(outputs),
+        }
+    finally:
+        client.stop_channels()
+        del client
+
+
+def _output_hook(ycell, outputs, msg) -> None:
+    msg_type = msg["header"]["msg_type"]
+    if msg_type in ("display_data", "stream", "execute_result", "error"):
+        # FIXME support for version
+        output = nbformat.v4.output_from_msg(msg)
+        outputs.append(output)
+
+        if ycell is not None:
+            cell_outputs = ycell["outputs"]
+            if msg_type == "stream":
+                with cell_outputs.doc.transaction():
+                    text = output["text"]
+
+                    # FIXME Logic is quite complex at https://github.com/jupyterlab/jupyterlab/blob/7ae2d436fc410b0cff51042a3350ba71f54f4445/packages/outputarea/src/model.ts#L518
+                    if text.endswith((os.linesep, "\n")):
+                        text = text[:-1]
+
+                    if (not cell_outputs) or (cell_outputs[-1]["name"] != output["name"]):
+                        output["text"] = [text]
+                        cell_outputs.append(output)
+                    else:
+                        last_output = cell_outputs[-1]
+                        last_output["text"].append(text)
+                        cell_outputs[-1] = last_output
+            else:
+                with cell_outputs.doc.transaction():
+                    cell_outputs.append(output)
+
+    elif msg_type == "clear_output":
+        # FIXME msg.content.wait - if true should clear at the next message
+        outputs.clear()
+
+        if ycell is not None:
+            del ycell["outputs"][:]
+
+    elif msg_type == "update_display_data":
+        # FIXME
+        ...
+
+
 class ExecuteHandler(ExtensionHandlerMixin, APIHandler):
+    """Handle request for snippet execution."""
+
     def initialize(
         self,
         name: str,
-        ydoc_extension: "jupyter_server_ydoc.app.YDocExtension" | None,
+        ydoc_extension: jupyter_server_ydoc.app.YDocExtension | None,
+        execution_stack: ExecutionStack,
         *args: t.Any,
         **kwargs: t.Any,
     ) -> None:
         super().initialize(name, *args, **kwargs)
+        self._execution_stack = execution_stack
         self._ydoc = ydoc_extension
-        self._outputs = []
-        self._ycell: y.Map | None = None
 
     @tornado.web.authenticated
     async def post(self, kernel_id: str) -> None:
@@ -54,6 +273,7 @@ class ExecuteHandler(ExtensionHandlerMixin, APIHandler):
         body = self.get_json_body()
 
         snippet = body.get("code")
+        ycell = None
         # From RTC model
         if snippet is None:
             document_id = body.get("document_id")
@@ -83,7 +303,7 @@ class ExecuteHandler(ExtensionHandlerMixin, APIHandler):
 
             ycells = filter(lambda c: c["id"] == cell_id, notebook.ycells)
             try:
-                self._ycell = next(ycells)
+                ycell = next(ycells)
             except StopIteration:
                 msg = f"Cell with ID {cell_id} not found in document {document_id}."
                 get_logger().error(msg)
@@ -95,7 +315,7 @@ class ExecuteHandler(ExtensionHandlerMixin, APIHandler):
                 except StopIteration:
                     get_logger().warning("Multiple cells have the same ID '%s'.", cell_id)
 
-            if self._ycell["cell_type"] != "code":
+            if ycell["cell_type"] != "code":
                 msg = f"Cell with ID {cell_id} of document {document_id} is not of type code."
                 get_logger().error(msg)
                 raise tornado.web.HTTPError(
@@ -103,7 +323,7 @@ class ExecuteHandler(ExtensionHandlerMixin, APIHandler):
                     reason=msg,
                 )
 
-            snippet = str(self._ycell["source"])
+            snippet = str(ycell["source"])
 
         try:
             km = self.kernel_manager.get_kernel(kernel_id)
@@ -112,57 +332,117 @@ class ExecuteHandler(ExtensionHandlerMixin, APIHandler):
             get_logger().error(msg, exc_info=e)
             raise tornado.web.HTTPError(status_code=HTTPStatus.NOT_FOUND, reason=msg) from e
 
-        client = km.client()
+        uid = self._execution_stack.put(km, snippet, ycell)
 
-        if self._ycell is not None:
-            # Reset cell
-            del self._ycell["outputs"][:]
-            self._ycell["execution_count"] = None
+        self.set_status(HTTPStatus.ACCEPTED)
+        self.set_header("Location", f"/api/kernels/{kernel_id}/requests/{uid}")
+        self.finish("{}")
 
-        # FIXME set the username of client.session to server user
-        # FIXME we don't check if the session is consistent (aka the kernel is linked to the document) - should we?
-        try:
-            reply = await ensure_async(
-                client.execute_interactive(
-                    snippet,
-                    output_hook=self._output_hook,
-                    stdin_hook=self._stdin_hook if client.allow_stdin else None,
-                )
-            )
 
-            reply_content = reply["content"]
+class InputHandler(ExtensionHandlerMixin, APIHandler):
+    """Handle request for input reply."""
 
-            if self._ycell is not None:
-                self._ycell["execution_count"] = reply_content["execution_count"]
+    @tornado.web.authenticated
+    async def post(self, kernel_id: str) -> None:
+        """
+        Send an input value to kernel ``kernel_id``.
 
-            self.finish(
-                {
-                    "status": reply_content["status"],
-                    "execution_count": reply_content["execution_count"],
-                    # FIXME quid for buffers
-                    "outputs": json.dumps(self._outputs),
-                }
-            )
+        Args:
+            kernel_id: Kernel identifier
+
+        Json Body Required:
+            input (str): Input value
+        """
+
+        try:
+            km = self.kernel_manager.get_kernel(kernel_id)
+        except KeyError as e:
+            msg = f"Unknown kernel with id: {kernel_id}"
+            get_logger().error(msg, exc_info=e)
+            raise tornado.web.HTTPError(status_code=HTTPStatus.NOT_FOUND, reason=msg) from e
+
+        body = self.get_json_body()
+        client = km.client()
 
+        try:
+            # only send stdin reply if there *was not* another request
+            # or execution finished while we were reading.
+            if not (
+                await client.stdin_channel.msg_ready() or await client.shell_channel.msg_ready()
+            ):
+                client.input(body["input"])
         finally:
-            self._outputs.clear()
-            self._ycell = None
+            client.stop_channels()
             del client
 
-    def _output_hook(self, msg) -> None:
-        msg_type = msg["header"]["msg_type"]
-        if msg_type in ("display_data", "stream", "execute_result", "error"):
-            # FIXME support for version
-            output = nbformat.v4.output_from_msg(msg)
-            get_logger().info("Got an output. %s", output)
-            self._outputs.append(output)
+        self.set_status(HTTPStatus.CREATED)
+
+
+class RequestHandler(ExtensionHandlerMixin, APIHandler):
+    """Handler for /api/kernels/<kernel_id>/requests/<request_id>"""
+
+    def initialize(
+        self, name: str, execution_stack: ExecutionStack, *args: t.Any, **kwargs: t.Any
+    ) -> None:
+        super().initialize(name, *args, **kwargs)
+        self._stack = execution_stack
+
+    @tornado.web.authenticated
+    def get(self, kernel_id: str, request_id: str) -> None:
+        """`GET /api/kernels/<kernel_id>/requests/<request_id>` Returns the request ``uid`` status.
+
+        Status are:
+
+        * 200: Request result is returned
+        * 202: Request is pending
+        * 300: Request has a pending input
+        * 500: Request ends with errors
+
+        Args:
+            kernel_id: Kernel identifier
+            request_id: Request identifier
+
+        Raises:
+            404 if request ``uid`` does not exist
+        """
+        try:
+            r = self._stack.get(kernel_id, request_id)
+        except ValueError as err:
+            raise tornado.web.HTTPError(404, reason=str(err)) from err
+        else:
+            if r is None:
+                self.set_status(202)
+                self.finish("{}")
+            else:
+                if "error" in r:
+                    self.set_status(500)
+                    self.log.debug(f"{r}")
+                elif "input_request" in r:
+                    self.set_status(300)
+                    self.set_header("Location", f"/api/kernels/{kernel_id}/input")
+                else:
+                    self.set_status(200)
+                self.finish(json.dumps(r))
+
+    # FIXME
+    # @tornado.web.authenticated
+    # def delete(self, kernel_id: str, request_id: str) -> None:
+    #     """`DELETE /api/kernels/<kernel_id>/requests/<request_id>` cancels the request ``uid``.
 
-            if self._ycell is not None:
-                # FIXME support for 'stream'
-                outputs = self._ycell["outputs"]
-                with outputs.doc.transaction():
-                    outputs.append(output)
+    #     Status are:
+    #     * 204: Request cancelled
 
+    #     Args:
+    #         kernel_id: Kernel identifier
+    #         request_id: Request identifier
 
-    def _stdin_hook(self, msg) -> None:
-        get_logger().info("Code snippet execution is waiting for an input.")
+    #     Raises:
+    #         404 if request ``uid`` does not exist
+    #     """
+    #     try:
+    #         self._stack.cancel(kernel_id, request_id)
+    #     except ValueError as err:
+    #         raise tornado.web.HTTPError(404, reason=str(err)) from err
+    #     else:
+    #         self.set_status(204)
+    #         self.finish()
diff --git a/jupyter_server_nbmodel/log.py b/jupyter_server_nbmodel/log.py
index f4ae28f..b83ced4 100644
--- a/jupyter_server_nbmodel/log.py
+++ b/jupyter_server_nbmodel/log.py
@@ -1,21 +1,18 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional
 
 from traitlets.config import Application
 
 
 class _ExtensionLogger:
-    _LOGGER: Optional[logging.Logger] = None
+    _LOGGER: logging.Logger | None = None
 
     @classmethod
     def get_logger(cls) -> logging.Logger:
         if cls._LOGGER is None:
             app = Application.instance()
-            cls._LOGGER = logging.getLogger(
-                "{!s}.jupyter_server_nb_model".format(app.log.name)
-            )
+            cls._LOGGER = logging.getLogger(f"{app.log.name!s}.jupyter_server_nb_model")
             Application.clear_instance()
 
         return cls._LOGGER


Test patch:
diff --git a/conftest.py b/conftest.py
index c2b2eba..49f6d64 100644
--- a/conftest.py
+++ b/conftest.py
@@ -2,6 +2,11 @@ import pytest
 
 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}
+        }
+    }
diff --git a/jupyter_server_nbmodel/tests/test_handlers.py b/jupyter_server_nbmodel/tests/test_handlers.py
index f9e5426..511ca17 100644
--- a/jupyter_server_nbmodel/tests/test_handlers.py
+++ b/jupyter_server_nbmodel/tests/test_handlers.py
@@ -1,11 +1,42 @@
 import asyncio
+import datetime
 import json
-
-from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
+import re
 
 import pytest
+from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
 
 TEST_TIMEOUT = 60
+SLEEP = 0.1
+
+
+REQUEST_REGEX = re.compile(r"^/api/kernels/\w+-\w+-\w+-\w+-\w+/requests/\w+-\w+-\w+-\w+-\w+$")
+
+
+async def _wait_request(fetch, endpoint: str):
+    """Poll periodically to fetch the execution request result."""
+    start_time = datetime.datetime.now()
+
+    while (datetime.datetime.now() - start_time).total_seconds() < 0.9 * TEST_TIMEOUT:
+        await asyncio.sleep(SLEEP)
+        response = await fetch(endpoint, raise_error=False)
+        if response.code >= 400:
+            response.rethrow()
+        if response.code != 202:
+            return response
+
+    raise TimeoutError(f"Request {endpoint} timed out.")
+
+
+async def wait_for_request(fetch, *args, **kwargs):
+    """Wait for execution request."""
+    r = await fetch(*args, **kwargs)
+    assert r.code == 202
+    location = r.headers["Location"]
+    assert REQUEST_REGEX.match(location) is not None
+
+    ans = await _wait_request(fetch, location)
+    return ans
 
 
 @pytest.fixture()
@@ -39,7 +70,7 @@ def pending_kernel_is_ready(jp_serverapp):
         (
             """from IPython.display import HTML
 HTML('<p><b>Jupyter</b> rocks.</p>')""",
-            '{"output_type": "execute_result", "metadata": {}, "data": {"text/plain": "<IPython.core.display.HTML object>", "text/html": "<p><b>Jupyter</b> rocks.</p>"}, "execution_count": 1}',
+            '{"output_type": "execute_result", "metadata": {}, "data": {"text/plain": "<IPython.core.display.HTML object>", "text/html": "<p><b>Jupyter</b> rocks.</p>"}, "execution_count": 1}',  # noqa: E501
         ),
     ),
 )
@@ -50,7 +81,8 @@ async def test_post_execute(jp_fetch, pending_kernel_is_ready, snippet, output):
     kernel = json.loads(r.body.decode())
     await pending_kernel_is_ready(kernel["id"])
 
-    response = await jp_fetch(
+    response = await wait_for_request(
+        jp_fetch,
         "api",
         "kernels",
         kernel["id"],
@@ -89,7 +121,8 @@ async def test_post_erroneous_execute(jp_fetch, pending_kernel_is_ready, snippet
     kernel = json.loads(r.body.decode())
     await pending_kernel_is_ready(kernel["id"])
 
-    response = await jp_fetch(
+    response = await wait_for_request(
+        jp_fetch,
         "api",
         "kernels",
         kernel["id"],
@@ -108,3 +141,92 @@ async def test_post_erroneous_execute(jp_fetch, pending_kernel_is_ready, snippet
 
     response2 = await jp_fetch("api", "kernels", kernel["id"], method="DELETE")
     assert response2.code == 204
+
+
+@pytest.mark.timeout(TEST_TIMEOUT)
+async def test_post_input_execute(jp_fetch, pending_kernel_is_ready):
+    # Start the first kernel
+    r = await jp_fetch(
+        "api", "kernels", method="POST", body=json.dumps({"name": NATIVE_KERNEL_NAME})
+    )
+    kernel = json.loads(r.body.decode())
+    await pending_kernel_is_ready(kernel["id"])
+
+    response = await jp_fetch(
+        "api",
+        "kernels",
+        kernel["id"],
+        "execute",
+        method="POST",
+        body=json.dumps({"code": "input('Age:')"}),
+    )
+    assert response.code == 202
+    location = response.headers["Location"]
+
+    response2 = await _wait_request(jp_fetch, location)
+
+    assert response2.code == 300
+    payload = json.loads(response2.body)
+    assert "parent_header" in payload
+    assert payload["input_request"] == {"prompt": "Age:", "password": False}
+
+    response3 = await jp_fetch(
+        "api", "kernels", kernel["id"], "input", method="POST", body=json.dumps({"input": "42"})
+    )
+    assert response3.code == 201
+
+    response4 = await _wait_request(
+        jp_fetch,
+        location
+    )
+    assert response4.code == 200
+    payload2 = json.loads(response4.body)
+    assert payload2 == {
+        "status": "ok",
+        "execution_count": 1,
+        "outputs": '[{"output_type": "execute_result", "metadata": {}, "data": {"text/plain": "\'42\'"}, "execution_count": 1}]',
+    }
+
+    r2 = await jp_fetch("api", "kernels", kernel["id"], method="DELETE")
+    assert r2.code == 204
+
+
+# FIXME
+# @pytest.mark.timeout(TEST_TIMEOUT)
+# async def test_cancel_execute(jp_fetch, pending_kernel_is_ready):
+#     # Start the first kernel
+#     r = await jp_fetch(
+#         "api", "kernels", method="POST", body=json.dumps({"name": NATIVE_KERNEL_NAME})
+#     )
+#     kernel = json.loads(r.body.decode())
+#     await pending_kernel_is_ready(kernel["id"])
+
+#     response = await jp_fetch(
+#         "api",
+#         "kernels",
+#         kernel["id"],
+#         "execute",
+#         method="POST",
+#         body=json.dumps({"code": """import time
+# time.sleep(10)
+# print("end")
+# """}),
+#     )
+
+#     assert response.code == 202
+#     location = response.headers["Location"]
+
+#     # Cancel task
+#     response2 = await jp_fetch(location, method="DELETE")
+
+#     assert response2.code == 204
+
+#     response3 = await jp_fetch(location)
+#     payload = json.loads(response3.body)
+#     assert payload == {
+#         "status": "error",
+#         "execution_count": 1
+#     }
+
+#     r2 = await jp_fetch("api", "kernels", kernel["id"], method="DELETE")
+#     assert r2.code == 204