Implement opt-in countdown feature and add to image recipe

stack/lab/Dockerfile

@@ -45,9 +45,25 @@ RUN jupyter nbextension enable --py --sys-prefix appmode && \
      jupyter serverextension enable --py --sys-prefix appmode
 
 # Swap appmode icon for AiiDAlab gears icon, shown during app load
-COPY --chown=${NB_UID}:${NB_GID} gears.svg ${CONDA_DIR}/share/jupyter/nbextensions/appmode/gears.svg
+COPY gears.svg ${CONDA_DIR}/share/jupyter/nbextensions/appmode/gears.svg
+
+ARG PYTHON_MINOR_VERSION
+
+# Set up the logo of notebook interface
+COPY aiidalab-wide-logo.png ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/base/images/logo.png
+
+# Copy custom CSS and JS files for AiiDAlab container countdown feature
+COPY countdown/ ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/custom/
+
+# Add endpoint to fetch container uptime
+COPY aiidalab_uptime.py ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/
+COPY aiidalab_uptime.json ${CONDA_DIR}/etc/jupyter/jupyter_notebook_config.d/
 
 # Copy start-up scripts for AiiDAlab.
+# We expose PYTHON_MINOR_VERSION in the container, as it is used by the
+# before-notebook.d/70_prepare_countdown_config.sh script to write the
+# environment-variable-dependent (opt-in) configuration file.
+ENV PYTHON_MINOR_VERSION=${PYTHON_MINOR_VERSION}
 COPY before-notebook.d/* /usr/local/bin/before-notebook.d/
 
 # Configure AiiDAlab environment.
@@ -71,12 +87,6 @@ ENV AIIDALAB_DEFAULT_APPS=""
 # Specify default factory reset (not set):
 ENV AIIDALAB_FACTORY_RESET=""
 
-USER ${NB_USER}
-
-WORKDIR "/home/${NB_USER}"
-
-RUN  mkdir -p /home/${NB_USER}/apps
-
 # When a Jupyter notebook server looses a connection to the frontend,
 # it keeps the messages in a buffer. If there is a background thread running
 # and trying to update the frontend, the buffer grows indefinitely,
@@ -98,6 +108,8 @@ ENV NOTEBOOK_ARGS=\
 "--TerminalManager.cull_inactive_timeout=3600 "\
 "--TerminalManager.cull_interval=300"
 
-# Set up the logo of notebook interface
-ARG PYTHON_MINOR_VERSION
-COPY --chown=${NB_UID}:${NB_GID} aiidalab-wide-logo.png ${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/base/images/logo.png
+USER ${NB_USER}
+
+WORKDIR "/home/${NB_USER}"
+
+RUN  mkdir -p /home/${NB_USER}/apps

stack/lab/aiidalab_uptime.json

@@ -0,0 +1,7 @@
+{
+  "NotebookApp": {
+    "nbserver_extensions": {
+      "aiidalab_uptime": true
+    }
+  }
+}

stack/lab/aiidalab_uptime.py

@@ -0,0 +1,32 @@
+import re
+import subprocess
+
+from notebook.base.handlers import IPythonHandler
+from notebook.utils import url_path_join
+
+
+class UptimeHandler(IPythonHandler):
+    def get(self):
+        try:
+            output = subprocess.check_output(["ps", "-p", "1", "-o", "etime="])
+            etime = output.decode("utf-8").strip()
+            seconds = self._parse_etime_to_seconds(etime)
+            self.finish({"uptime": seconds})
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": str(e)})
+
+    def _parse_etime_to_seconds(self, etime):
+        # Supports MM:SS, HH:MM:SS, or D-HH:MM:SS formats
+        match = re.match(r"(?:(\d+)-)?(?:(\d+):)?(\d+):(\d+)", etime)
+        if not match:
+            raise ValueError(f"Unrecognized etime format: {etime}")
+
+        days, hours, minutes, seconds = match.groups(default="0")
+        return int(days) * 86400 + int(hours) * 3600 + int(minutes) * 60 + int(seconds)
+
+
+def load_jupyter_server_extension(nb_server_app):
+    web_app = nb_server_app.web_app
+    route = url_path_join(web_app.settings["base_url"], "/uptime")
+    web_app.add_handlers(".*", [(route, UptimeHandler)])

stack/lab/before-notebook.d/70_prepare_countdown_config.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+set -e
+
+CUSTOM_DIR="${CONDA_DIR}/lib/python${PYTHON_MINOR_VERSION}/site-packages/notebook/static/custom"
+
+cat <<EOF >"${CUSTOM_DIR}/config.json"
+{
+  "ephemeral": $([ -n "$LIFETIME" ] && echo 1 || echo 0),
+  "lifetime": "$LIFETIME"
+}
+EOF

stack/lab/countdown/custom.css

@@ -0,0 +1,33 @@
+#culling-countdown {
+  position: sticky;
+  top: 0;
+  left: 0;
+  width: 100%;
+  background-color: #0078d4;
+  color: white;
+  text-align: center;
+  padding: 8px;
+  font-size: 18px;
+  font-weight: bold;
+  z-index: 9999;
+}
+
+#shutdown-warning,
+#save-info {
+  display: none;
+}
+
+#shutdown-warning {
+  font-size: 20px;
+}
+
+#save-info {
+  font-size: 16px;
+  text-align: center;
+  font-weight: normal;
+  font-size: 18px;
+}
+
+#culling-timer {
+  margin-left: 5px;
+}

stack/lab/countdown/custom.js

@@ -0,0 +1,116 @@
+require(["base/js/namespace", "base/js/events"], (Jupyter, events) => {
+  const parseLifetimeToMs = (str) => {
+    // Supports MM:SS, HH:MM:SS, or D-HH:MM:SS formats
+    const regex = /^(?:(\d+)-)?(?:(\d+):)?(\d+):(\d+)$/;
+    const match = regex.exec(str.trim());
+
+    if (!match) {
+      console.warn("Invalid lifetime format:", str);
+      return null;
+    }
+
+    const [_, days, hours, minutes, seconds] = match.map((v) => Number(v) || 0);
+
+    const totalSeconds = days * 86400 + hours * 3600 + minutes * 60 + seconds;
+    return totalSeconds * 1000;
+  };
+
+  const insertCountdown = (remainingMs) => {
+    if (document.getElementById("culling-countdown")) {
+      return;
+    }
+
+    const banner = document.createElement("div");
+    banner.id = "culling-countdown";
+
+    const shutdownWarning = document.createElement("div");
+    shutdownWarning.id = "shutdown-warning";
+    shutdownWarning.innerHTML = "⚠️ Shutdown imminent! ⚠️";
+    banner.appendChild(shutdownWarning);
+
+    const countdown = document.createElement("div");
+    countdown.id = "countdown";
+    countdown.innerHTML = `Session time remaining: `;
+    const timer = document.createElement("span");
+    timer.id = "culling-timer";
+    timer.innerHTML = "Calculating...";
+    countdown.appendChild(timer);
+    banner.appendChild(countdown);
+
+    const saveInfo = document.createElement("div");
+    saveInfo.id = "save-info";
+    saveInfo.innerHTML = `
+      Consider saving your work using the <b>File Manager</b> or the <b>Terminal</b>
+    `;
+    banner.appendChild(saveInfo);
+
+    const startTime = Date.now();
+    const endTime = startTime + remainingMs;
+
+    const formatTime = (seconds) => {
+      const hrs = `${Math.floor(seconds / 3600)}`.padStart(2, "0");
+      const mins = `${Math.floor((seconds % 3600) / 60)}`.padStart(2, "0");
+      const secs = `${Math.floor(seconds % 60)}`.padStart(2, "0");
+      // Format as HH:MM:SS, even if > 1 day (for now - rare?)
+      return `${hrs}:${mins}:${secs}`;
+    };
+
+    const updateTimer = () => {
+      const timeLeft = (endTime - Date.now()) / 1000;
+      if (timeLeft < 0) {
+        clearInterval(interval);
+        return;
+      }
+      if (timeLeft < 1800) {
+        banner.style.backgroundColor = "#DAA801";
+        saveInfo.style.display = "block";
+      }
+      if (timeLeft < 300) {
+        banner.style.backgroundColor = "red";
+        shutdownWarning.style.display = "block";
+        shutdownWarning.innerHTML = "⚠️ Shutdown imminent ⚠️";
+      }
+      timer.innerHTML = formatTime(timeLeft);
+    };
+
+    updateTimer();
+    const interval = setInterval(updateTimer, 1000);
+
+    const container = document.getElementById("header");
+    if (container) {
+      container.parentNode.insertBefore(banner, container);
+    }
+  };
+
+  const loadCountdown = async () => {
+    try {
+      const [configResponse, uptimeResponse] = await Promise.all([
+        fetch("/static/custom/config.json"),
+        fetch("/uptime"),
+      ]);
+
+      const config = await configResponse.json();
+      const uptimeData = await uptimeResponse.json();
+
+      if (!config.ephemeral) {
+        return;
+      }
+
+      if (!config.lifetime) {
+        console.warn("Missing `lifetime` in config file");
+        return;
+      }
+
+      const lifetimeMs = parseLifetimeToMs(config.lifetime);
+      const uptimeMs = uptimeData.uptime * 1000;
+      const remaining = lifetimeMs - uptimeMs;
+
+      insertCountdown(Math.max(0, remaining));
+    } catch (err) {
+      console.error("Countdown init failed:", err);
+    }
+  };
+
+  events.on("app_initialized.NotebookApp", loadCountdown);
+  loadCountdown();
+});