docker: address #6681 review round 2 (colab magics, output select, branding guard)

- unsloth_colab_compat.py: only hoist a leading `%%` cell magic above the Colab
  `#@title` form for magics whose body runs as code (capture/time/bash/python/
  ...). Content magics (%%writefile, %%html, %%latex, ...) are left untouched so
  the form comment is never injected into the written file / rendered output.
- outputSelect.ts: stop trusting the text selection anchor to decide ownership
  of Ctrl/Cmd+A. A stale selection inside an output survives a click onto a
  command-mode cell or the file browser, which made select-all keep re-selecting
  the old output. Gate on the keystroke target or the last pointer-down (reset to
  null on any click outside an output) instead.
- unsloth_branding.py: also reject page_config.json that disables the Unsloth
  labextension or any of its plugin ids via disabledExtensions (dict or list
  form); that leaves the bundle on disk so the prior checks passed while the
  logo/About/splash attribution was stripped at load. Lock unsloth-jupyterlab in
  Dockerfile.studio as well (defense in depth), and add guard tests.

docker/Dockerfile.studio

@@ -240,7 +240,8 @@ RUN JS="$(/opt/unsloth-venv/bin/python -c 'import os, jupyter_server; print(os.p
  && /opt/unsloth-venv/bin/jupyter labextension disable @jupyterlab/application-extension:logo \
  && /opt/unsloth-venv/bin/jupyter labextension lock @jupyterlab/application-extension:logo \
  && /opt/unsloth-venv/bin/jupyter labextension disable @jupyterlab/apputils-extension:splash \
- && /opt/unsloth-venv/bin/jupyter labextension lock @jupyterlab/apputils-extension:splash
+ && /opt/unsloth-venv/bin/jupyter labextension lock @jupyterlab/apputils-extension:splash \
+ && /opt/unsloth-venv/bin/jupyter labextension lock unsloth-jupyterlab
 # Branding integrity guard: the canonical attribution checker (also a
 # jupyter_server extension), the full AGPLv3 license text, and the config that
 # enables the extension. Installed into the base venv so it is on the jupyter

docker/jupyter/unsloth_branding.py

@@ -58,7 +58,7 @@ SPLASH_PLUGIN_ID = "unsloth-jupyterlab:splash"
 LOGO_DATA_URI_PREFIX = "data:image/png;base64,iVBOR"
 
 
-def resolve_paths(venv_share=None, jupyter_server_dir=None):
+def resolve_paths(venv_share=None, jupyter_server_dir=None, config_dirs=None):
     """Resolve the installed locations of every checked branding asset.
 
     Defaults point at the live venv + the installed jupyter_server package. Tests
@@ -71,6 +71,21 @@ def resolve_paths(venv_share=None, jupyter_server_dir=None):
 
         jupyter_server_dir = os.path.dirname(jupyter_server.__file__)
     labext_dir = os.path.join(venv_share, "labextensions", LABEXT_NAME)
+
+    # Every page_config.json JupyterLab merges to compute disabledExtensions: the
+    # app-settings file plus a labconfig/ file under each jupyter config dir
+    # (where `jupyter labextension disable` writes). Tests pass config_dirs=[] for
+    # a hermetic tree; live resolution scans the real jupyter config path.
+    if config_dirs is None:
+        try:
+            from jupyter_core.paths import jupyter_config_path
+
+            config_dirs = jupyter_config_path()
+        except Exception:
+            config_dirs = []
+    page_configs = [os.path.join(venv_share, "lab", "settings", "page_config.json")]
+    page_configs += [os.path.join(d, "labconfig", "page_config.json") for d in config_dirs]
+
     return {
         "license": os.path.join(venv_share, "UNSLOTH_LICENSE.AGPL-3.0"),
         "login": os.path.join(jupyter_server_dir, "templates", "login.html"),
@@ -80,6 +95,7 @@ def resolve_paths(venv_share=None, jupyter_server_dir=None):
         "labext_static": os.path.join(labext_dir, "static"),
         "favicon": os.path.join(jupyter_server_dir, "static", "favicons", "favicon.ico"),
         "logo": os.path.join(jupyter_server_dir, "static", "logo", "logo.png"),
+        "page_configs": page_configs,
     }
 
 
@@ -177,6 +193,36 @@ def verify_branding(paths = None):
     if not _nonempty_file(paths["logo"]):
         problems.append("missing or empty logo: " + paths["logo"])
 
+    # 7. No page_config.json disables the Unsloth extension or its plugins.
+    #    Disabling via `disabledExtensions` leaves the static bundle on disk (so
+    #    check 5 still passes) yet strips the logo / About / splash at load. Since
+    #    the guard exists to refuse stripped attribution, reject that too. Stock
+    #    plugins we disable ourselves (logo/splash) are unaffected -- we only flag
+    #    ids belonging to unsloth-jupyterlab.
+    for pc_path in paths.get("page_configs", []):
+        text = _read(pc_path)
+        if not text:
+            continue
+        try:
+            disabled = json.loads(text).get("disabledExtensions", {})
+        except ValueError:
+            problems.append("page_config.json is not valid JSON: " + pc_path)
+            continue
+        # Modern JupyterLab uses a {id: bool} map; older configs used a list.
+        if isinstance(disabled, dict):
+            disabled_ids = [k for k, v in disabled.items() if v]
+        elif isinstance(disabled, (list, tuple)):
+            disabled_ids = list(disabled)
+        else:
+            disabled_ids = []
+        for ident in disabled_ids:
+            if not isinstance(ident, str):
+                continue
+            if ident == LABEXT_NAME or ident.startswith(LABEXT_NAME + ":"):
+                problems.append(
+                    "page_config.json disables Unsloth attribution '" + ident + "': " + pc_path
+                )
+
     return problems
 
 

docker/jupyter/unsloth_labext/src/outputSelect.ts

@@ -22,9 +22,16 @@ import {
  *   - the chord is exactly Ctrl/Cmd+A (no Alt; Shift ignored), and
  *   - focus is NOT in a text editor / input / contenteditable (so editing a
  *     code cell with Ctrl+A still selects within that editor), and
- *   - the current selection anchor or the last pointer-down landed inside an
- *     output area.
- * In every other case we do nothing and JupyterLab keeps its default behaviour.
+ *   - the keystroke target OR the last pointer-down landed inside an output area.
+ *
+ * We deliberately do NOT use the text selection anchor to decide ownership: a
+ * stale selection inside an output survives a later click onto a command-mode
+ * cell or the file browser (clicking a non-text region does not always move the
+ * anchor), which would make Ctrl/Cmd+A keep re-selecting that old output instead
+ * of doing the normal select-all in the new context. The last pointer-down is
+ * reset on every click (to null when the click is outside any output), so it
+ * tracks the user's current intent; in every other case we do nothing and
+ * JupyterLab keeps its default behaviour.
  */
 
 // Output containers, widest first. `.jp-OutputArea-output` is a single output;
@@ -95,13 +102,13 @@ const outputSelectPlugin: JupyterFrontEndPlugin<void> = {
       if (inEditableContext()) {
         return;
       }
-      // Prefer the output holding the current selection; fall back to the last
-      // place the user clicked.
-      const selection = window.getSelection();
-      let output = closestOutput(selection?.anchorNode ?? null);
-      if (!output) {
-        output = lastPointerOutput;
-      }
+      // Own the chord only when the user is actually in an output right now:
+      // the keystroke target, else the last place they clicked. We do NOT trust
+      // the text selection anchor -- it goes stale after clicking away from a
+      // previously selected output (see the file header), which would otherwise
+      // hijack select-all in the notebook / file browser.
+      const output =
+        closestOutput(event.target as Node | null) ?? lastPointerOutput;
       if (!output) {
         return;
       }

docker/unsloth_colab_compat.py

@@ -19,6 +19,13 @@ stay in the cell (still inert), just below the magic -- so `%%capture` now also
 captures them. Idempotent and fully guarded: any problem returns the input
 unchanged, so a cell never breaks because of this helper.
 
+The hoist is restricted to cell magics whose body is executed as code (Python or
+shell), where a moved-down `#@title`/comment line stays an inert comment. Magics
+that treat the body as literal content (`%%writefile`, `%%file`, `%%html`,
+`%%javascript`, `%%latex`, `%%markdown`, `%%svg`, ...) are left untouched: moving
+the Colab form comment into their body would write/render it and corrupt the
+generated file or output.
+
 This mirrors unsloth_nb_compat.register_ipython(): it is wired from the baked
 IPython startup file (docker/unsloth_ipython_startup.py).
 """
@@ -27,8 +34,21 @@ from __future__ import annotations
 import sys
 
 
+# Cell magics whose body is executed as code (Python or shell), so a hoisted
+# `#@title`/`#@param`/comment line stays an inert comment. We ONLY hoist these.
+# Anything not listed (content/data magics like %%writefile, %%file, %%html,
+# %%javascript, %%latex, %%markdown, %%svg) is left untouched, because injecting
+# the Colab form comment into its body would corrupt the written file / output.
+_SAFE_CELL_MAGICS = frozenset({
+    "capture",  # the Colab install pattern: suppress pip/install output
+    "time", "timeit", "prun", "debug",
+    "bash", "sh", "shell",
+    "python", "python2", "python3", "pypy",
+})
+
+
 def colab_cell_magic_fix(lines):
-    """Hoist a `%%` cell magic above leading blank/comment lines.
+    """Hoist a safe `%%` cell magic above leading blank/comment lines.
 
     `lines` is the IPython cell as a list of strings (each ending in '\\n').
     Returns a (possibly reordered) list of the same lines.
@@ -43,7 +63,13 @@ def colab_cell_magic_fix(lines):
             # First real line. Only act if it is a cell magic that is not yet on
             # top (i.e. something was skipped before it).
             if stripped.startswith("%%") and i > 0:
-                return [line] + skipped + lines[i + 1 :]
+                name = stripped[2:].split(maxsplit=1)
+                name = name[0] if name else ""
+                if name in _SAFE_CELL_MAGICS:
+                    return [line] + skipped + lines[i + 1 :]
+                # Content/data magic (%%writefile, %%html, ...): do not move the
+                # comment into its body. Leave the cell exactly as written.
+                return lines
             return lines  # already on top, or not a magic
         return lines  # all blank/comment -> nothing to do
     except Exception:

tests/studio/test_branding_guard.py

@@ -69,7 +69,11 @@ def _stage(tmp_path):
     (js_dir / "static" / "logo").mkdir(parents = True)
     (js_dir / "static" / "logo" / "logo.png").write_bytes(b"\x89PNG\r\n\x1a\nlogo")
 
-    return ub.resolve_paths(venv_share = str(venv_share), jupyter_server_dir = str(js_dir))
+    # config_dirs = [] keeps the tree hermetic (no host jupyter config scanned);
+    # page_config tests write to the app-settings page_config.json directly.
+    return ub.resolve_paths(
+        venv_share = str(venv_share), jupyter_server_dir = str(js_dir), config_dirs = [],
+    )
 
 
 def test_positive_clean_tree_passes(tmp_path):
@@ -143,6 +147,22 @@ def _empty_favicon(paths):
     open(paths["favicon"], "w").close()
 
 
+def _disable_unsloth_ext(paths):
+    with open(paths["page_configs"][0], "w", encoding = "utf-8") as f:
+        json.dump({"disabledExtensions": {ub.LABEXT_NAME: True}}, f)
+
+
+def _disable_unsloth_plugin(paths):
+    with open(paths["page_configs"][0], "w", encoding = "utf-8") as f:
+        json.dump({"disabledExtensions": {ub.ABOUT_PLUGIN_ID: True}}, f)
+
+
+def _disable_unsloth_ext_list_form(paths):
+    # Older JupyterLab configs used a list of ids rather than an {id: bool} map.
+    with open(paths["page_configs"][0], "w", encoding = "utf-8") as f:
+        json.dump({"disabledExtensions": [ub.SPLASH_PLUGIN_ID]}, f)
+
+
 @pytest.mark.parametrize(
     "mutate",
     [
@@ -157,6 +177,9 @@ def _empty_favicon(paths):
         _strip_bundle_logo,
         _remove_logo_png,
         _empty_favicon,
+        _disable_unsloth_ext,
+        _disable_unsloth_plugin,
+        _disable_unsloth_ext_list_form,
     ],
 )
 def test_negative_each_marker_is_enforced(tmp_path, mutate):
@@ -167,6 +190,17 @@ def test_negative_each_marker_is_enforced(tmp_path, mutate):
     assert problems, "stripping " + mutate.__name__ + " must be detected"
 
 
+def test_disabling_stock_plugins_is_allowed(tmp_path):
+    """We disable the stock logo/splash ourselves -- the guard must not flag those."""
+    paths = _stage(tmp_path)
+    with open(paths["page_configs"][0], "w", encoding = "utf-8") as f:
+        json.dump({"disabledExtensions": {
+            "@jupyterlab/application-extension:logo": True,
+            "@jupyterlab/apputils-extension:splash": True,
+        }}, f)
+    assert ub.verify_branding(paths) == []
+
+
 def test_attribution_sources_have_no_encoded_obfuscation():
     """Plain readable strings only -- no base64/decoder tricks (antivirus-safe)."""
     src_dir = os.path.join(REPO, "docker", "jupyter")

tests/validate_studio_features.py

@@ -53,6 +53,13 @@ def test_colab_compat() -> None:
     # non-magic cell untouched
     plain = ["x = 1\n", "y = 2\n"]
     check("plain cell untouched", m.colab_cell_magic_fix(plain) == plain)
+    # content/data magic (%%writefile) NOT hoisted -- never inject the #@title
+    # comment into the written file body
+    wf = ["#@title Config\n", "%%writefile config.json\n", "{}\n"]
+    check("content magic (%%writefile) left untouched", m.colab_cell_magic_fix(wf) == wf)
+    # safe magic with arg still hoisted
+    bash = ["#@title Run\n", "%%bash\n", "echo hi\n"]
+    check("safe magic (%%bash) hoisted", m.colab_cell_magic_fix(bash)[0] == "%%bash\n")
 
 
 # --------------------------------------------------------------------------