From 09684055c22261475c083ecb32863308246655e0 Mon Sep 17 00:00:00 2001
From: Jeremy Tuloup <jeremy.tuloup@gmail.com>
Date: Thu, 21 Aug 2025 09:17:11 +0200
Subject: [PATCH] Increase checkpoint poll interval

---
 .../schema/checkpoints.json                   |  10 +-
 packages/notebook-extension/src/index.ts      | 106 +++++++++++++++---
 2 files changed, 98 insertions(+), 18 deletions(-)

diff --git a/packages/notebook-extension/schema/checkpoints.json b/packages/notebook-extension/schema/checkpoints.json
index 7232782395..1d3f752549 100644
--- a/packages/notebook-extension/schema/checkpoints.json
+++ b/packages/notebook-extension/schema/checkpoints.json
@@ -4,7 +4,15 @@
   "jupyter.lab.toolbars": {
     "TopBar": [{ "name": "checkpoint", "rank": 20 }]
   },
-  "properties": {},
+  "properties": {
+    "checkpointPollingInterval": {
+      "type": "number",
+      "title": "Checkpoint Polling Interval (seconds)",
+      "description": "How often to check for checkpoints (in seconds). Set to 0 to disable polling.",
+      "default": 30,
+      "minimum": 0
+    }
+  },
   "additionalProperties": false,
   "type": "object"
 }
diff --git a/packages/notebook-extension/src/index.ts b/packages/notebook-extension/src/index.ts
index cd9d9752d1..e0fafbb08f 100644
--- a/packages/notebook-extension/src/index.ts
+++ b/packages/notebook-extension/src/index.ts
@@ -19,6 +19,8 @@ import { PageConfig, Text, Time, URLExt } from '@jupyterlab/coreutils';
 
 import { IDocumentManager } from '@jupyterlab/docmanager';
 
+import { DocumentRegistry } from '@jupyterlab/docregistry';
+
 import { IMainMenu } from '@jupyterlab/mainmenu';
 
 import {
@@ -92,13 +94,14 @@ const checkpoints: JupyterFrontEndPlugin<void> = {
   description: 'A plugin for the checkpoint indicator.',
   autoStart: true,
   requires: [IDocumentManager, ITranslator],
-  optional: [INotebookShell, IToolbarWidgetRegistry],
+  optional: [INotebookShell, IToolbarWidgetRegistry, ISettingRegistry],
   activate: (
     app: JupyterFrontEnd,
     docManager: IDocumentManager,
     translator: ITranslator,
     notebookShell: INotebookShell | null,
-    toolbarRegistry: IToolbarWidgetRegistry | null
+    toolbarRegistry: IToolbarWidgetRegistry | null,
+    settingRegistry: ISettingRegistry | null
   ) => {
     const { shell } = app;
     const trans = translator.load('notebook');
@@ -113,18 +116,26 @@ const checkpoints: JupyterFrontEndPlugin<void> = {
       });
     }
 
-    const onChange = async () => {
+    const getCurrent = () => {
       const current = shell.currentWidget;
       if (!current) {
-        return;
+        return null;
       }
       const context = docManager.contextForWidget(current);
+      if (!context) {
+        return null;
+      }
+      return context;
+    };
 
-      context?.fileChanged.disconnect(onChange);
-      context?.fileChanged.connect(onChange);
-
-      const checkpoints = await context?.listCheckpoints();
+    const updateCheckpointDisplay = async () => {
+      const current = getCurrent();
+      if (!current) {
+        return;
+      }
+      const checkpoints = await current.listCheckpoints();
       if (!checkpoints || !checkpoints.length) {
+        node.textContent = '';
         return;
       }
       const checkpoint = checkpoints[checkpoints.length - 1];
@@ -134,19 +145,80 @@ const checkpoints: JupyterFrontEndPlugin<void> = {
       );
     };
 
+    const onSaveState = async (
+      sender: DocumentRegistry.IContext<DocumentRegistry.IModel>,
+      state: DocumentRegistry.SaveState
+    ) => {
+      if (state !== 'completed') {
+        return;
+      }
+      // Add a small artificial delay so that the UI can pick up the newly created checkpoint.
+      // Since the save state signal is emitted after a file save, but not after a checkpoint has been created.
+      setTimeout(() => {
+        void updateCheckpointDisplay();
+      }, 500);
+    };
+
+    const onChange = async () => {
+      const context = getCurrent();
+      if (!context) {
+        return;
+      }
+
+      context.saveState.disconnect(onSaveState);
+      context.saveState.connect(onSaveState);
+
+      await updateCheckpointDisplay();
+    };
+
     if (notebookShell) {
       notebookShell.currentChanged.connect(onChange);
     }
 
-    new Poll({
-      auto: true,
-      factory: () => onChange(),
-      frequency: {
-        interval: 2000,
-        backoff: false,
-      },
-      standby: 'when-hidden',
-    });
+    let checkpointPollingInterval = 30; // Default 30 seconds
+    let poll: Poll | null = null;
+
+    const createPoll = () => {
+      if (poll) {
+        poll.dispose();
+      }
+      if (checkpointPollingInterval > 0) {
+        poll = new Poll({
+          auto: true,
+          factory: () => updateCheckpointDisplay(),
+          frequency: {
+            interval: checkpointPollingInterval * 1000,
+            backoff: false,
+          },
+          standby: 'when-hidden',
+        });
+      }
+    };
+
+    const updateSettings = (settings: ISettingRegistry.ISettings): void => {
+      checkpointPollingInterval = settings.get('checkpointPollingInterval')
+        .composite as number;
+      createPoll();
+    };
+
+    if (settingRegistry) {
+      const loadSettings = settingRegistry.load(checkpoints.id);
+      Promise.all([loadSettings, app.restored])
+        .then(([settings]) => {
+          updateSettings(settings);
+          settings.changed.connect(updateSettings);
+        })
+        .catch((reason: Error) => {
+          console.error(
+            `Failed to load settings for ${checkpoints.id}: ${reason.message}`
+          );
+          // Fall back to creating poll with default settings
+          createPoll();
+        });
+    } else {
+      // Create poll with default settings
+      createPoll();
+    }
   },
 };