Add execution start and end time metadata for code cells

conftest.py

@@ -1,12 +1,31 @@
 import pytest
 
-pytest_plugins = ("pytest_jupyter.jupyter_server",)
-
+pytest_plugins = [
+    "pytest_jupyter.jupyter_server",
+    "jupyter_server.pytest_plugin",
+    "jupyter_server_fileid.pytest_plugin",
+    "jupyter_server_ydoc.pytest_plugin"
+]
 
 @pytest.fixture
-def jp_server_config(jp_server_config):
+def jp_server_config(jp_root_dir,jp_server_config):
     return {
-        "ServerApp": {
-            "jpserver_extensions": {"jupyter_server_nbmodel": True, "jupyter_server_ydoc": False}
-        }
-    }
+        'ServerApp': {
+            'jpserver_extensions': {
+                'jupyter_server_ydoc': True,
+                'jupyter_server_fileid': True,
+                'jupyter_server_nbmodel': True
+            },
+            'token': '',
+            'password': '',
+            'disable_check_xsrf': True},
+            "SQLiteYStore": {"db_path": str(jp_root_dir.joinpath(".rtc_test.db"))},
+            "BaseFileIdManager": {
+                "root_dir": str(jp_root_dir),
+                "db_path": str(jp_root_dir.joinpath(".fid_test.db")),
+                "db_journal_mode": "OFF",
+            },
+            'YDocExtension': {
+                'document_save_delay': 1
+            }
+    } 

jupyter_server_nbmodel/handlers.py

@@ -8,6 +8,7 @@
 from dataclasses import asdict, dataclass
 from functools import partial
 from http import HTTPStatus
+from datetime import datetime, timezone
 
 import jupyter_server
 import jupyter_server.services
@@ -123,7 +124,6 @@ async def _get_ycell(
         raise KeyError(
             msg,
         )
-
     return ycell
 
 
@@ -219,6 +219,7 @@ async def _execute_snippet(
         The execution status and outputs.
     """
     ycell = None
+    time_info = {}
     if metadata is not None:
         ycell = await _get_ycell(ydoc, metadata)
         if ycell is not None:
@@ -227,7 +228,11 @@ async def _execute_snippet(
                 del ycell["outputs"][:]
                 ycell["execution_count"] = None
                 ycell["execution_state"] = "running"
-
+                if metadata["record_timing"]:
+                    time_info = ycell["metadata"].get("execution",{})
+                    time_info["shell.execute_reply.started"] = datetime.now(timezone.utc).isoformat()[:-6]
+                    ycell["metadata"]["execution"] = time_info
+
     outputs = []
 
     # FIXME we don't check if the session is consistent (aka the kernel is linked to the document)
@@ -247,7 +252,13 @@ async def _execute_snippet(
         with ycell.doc.transaction():
             ycell["execution_count"] = reply_content.get("execution_count")
             ycell["execution_state"] = "idle"
-
+            if metadata["record_timing"]:
+                end_time = datetime.now(timezone.utc).isoformat()[:-6]
+                if reply_content["status"] == "ok":
+                    time_info["shell.execute_reply"] = end_time
+                else:
+                    time_info["execution_failed"] = end_time
+                ycell["metadata"]["execution"] = time_info
     return {
         "status": reply_content["status"],
         "execution_count": reply_content.get("execution_count"),
@@ -524,9 +535,7 @@ async def post(self, kernel_id: str) -> None:
             msg = f"Unknown kernel with id: {kernel_id}"
             get_logger().error(msg)
             raise tornado.web.HTTPError(status_code=HTTPStatus.NOT_FOUND, reason=msg)
-
         uid = self._execution_stack.put(kernel_id, snippet, metadata)
-
         self.set_status(HTTPStatus.ACCEPTED)
         self.set_header("Location", f"/api/kernels/{kernel_id}/requests/{uid}")
         self.finish("{}")

jupyter_server_nbmodel/tests/test_handlers.py

@@ -2,6 +2,7 @@
 import datetime
 import json
 import re
+import nbformat
 
 import pytest
 from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
@@ -148,6 +149,58 @@ async def test_post_erroneous_execute(jp_fetch, pending_kernel_is_ready, snippet
 
     await asyncio.sleep(1)
 
+@pytest.mark.timeout(TEST_TIMEOUT)
+async def test_execution_timing_metadata(jp_fetch, pending_kernel_is_ready, rtc_create_notebook, jp_serverapp):
+    snippet = "a = 1"
+    nb = nbformat.v4.new_notebook(
+        cells=[nbformat.v4.new_code_cell(source=snippet, execution_count=1)]
+    )
+    nb_content = nbformat.writes(nb, version=4)
+    path, _ = await rtc_create_notebook("test.ipynb", nb_content, store=True)
+    collaboration = jp_serverapp.web_app.settings["jupyter_server_ydoc"]
+    fim = jp_serverapp.web_app.settings["file_id_manager"]
+    document = await collaboration.get_document(
+        path=path, content_type="notebook", file_format="json", copy=False
+    )
+    doc = document.get()
+    document_id = f'json:notebook:{fim.get_id("test.ipynb")}'
+    cell_id = doc["cells"][0].get("id")
+
+    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 wait_for_request(
+        jp_fetch,
+        "api",
+        "kernels",
+        kernel["id"],
+        "execute",
+        method="POST",
+        body=json.dumps({"code": snippet, "metadata":{"cell_id":cell_id,"document_id":document_id,"record_timing":True}}),
+    )
+    assert response.code == 200
+    cell_data = document.get()["cells"][0]
+    assert 'execution' in cell_data['metadata'], "'execution' does not exist in 'metadata'"
+
+    # Assert that start and end time exist in 'execution'
+    execution = cell_data['metadata']['execution']
+    assert 'shell.execute_reply.started' in execution, "'shell.execute_reply.started' does not exist in 'execution'"
+    assert 'shell.execute_reply' in execution, "'shell.execute_reply' does not exist in 'execution'"
+
+    started_time = execution['shell.execute_reply.started']
+    reply_time = execution['shell.execute_reply']
+
+    started_dt = datetime.datetime.fromisoformat(started_time)
+    reply_dt = datetime.datetime.fromisoformat(reply_time)
+
+    # Assert that reply_time is greater than started_time
+    assert reply_dt > started_dt, "The reply time is not greater than the started time."
+    response2 = await jp_fetch("api", "kernels", kernel["id"], method="DELETE")
+    assert response2.code == 204
+    await asyncio.sleep(1)
 
 @pytest.mark.timeout(TEST_TIMEOUT)
 async def test_post_input_execute(jp_fetch, pending_kernel_is_ready):

pyproject.toml

@@ -30,7 +30,14 @@ dynamic = ["version", "description", "authors", "urls", "keywords"]
 
 [project.optional-dependencies]
 lab = ["jupyterlab>=4.2.0", "jupyter-docprovider>=1.0.0b1", "jupyter-server-ydoc>=1.0.0b1"]
-test = ["pytest~=8.2", "pytest-cov", "pytest-jupyter[server]>=0.6", "pytest-timeout"]
+test = [
+  "pytest~=8.2",
+  "pytest-cov",
+  "pytest-jupyter[server]>=0.6",
+  "pytest-timeout",
+  "jupyter-server-ydoc[test]>=1.0.0b1",
+  "jupyter-server-fileid"
+]
 lint = ["mdformat>0.7", "mdformat-gfm>=0.3.5", "ruff>=0.4.0"]
 typing = ["mypy>=0.990"]
 

src/executor.ts

@@ -108,12 +108,17 @@ export class NotebookCellServerExecutor implements INotebookCellExecutor {
           const code = cell.model.sharedModel.getSource();
           const cellId = cell.model.sharedModel.getId();
           const documentId = notebook.sharedModel.getState('document_id');
+          const { recordTiming } = notebookConfig;
 
           const init = {
             method: 'POST',
             body: JSON.stringify({
               code,
-              metadata: { cell_id: cellId, document_id: documentId }
+              metadata: {
+                cell_id: cellId,
+                document_id: documentId,
+                record_timing: recordTiming
+              }
             })
           };
           onCellExecutionScheduled({ cell });