From 2482719695e8f35f06f25b107e2cbb888daee21c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 27 Jun 2026 08:55:04 +0000 Subject: [PATCH] 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 | 3 +- docker/jupyter/unsloth_branding.py | 48 ++++++++++++++++++- .../unsloth_labext/src/outputSelect.ts | 27 +++++++---- docker/unsloth_colab_compat.py | 30 +++++++++++- tests/studio/test_branding_guard.py | 36 +++++++++++++- tests/validate_studio_features.py | 7 +++ 6 files changed, 136 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile.studio b/docker/Dockerfile.studio index d0f8cfb510..d6480704f3 100644 --- a/docker/Dockerfile.studio +++ b/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 diff --git a/docker/jupyter/unsloth_branding.py b/docker/jupyter/unsloth_branding.py index 44722dee7c..5287869a6a 100644 --- a/docker/jupyter/unsloth_branding.py +++ b/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 diff --git a/docker/jupyter/unsloth_labext/src/outputSelect.ts b/docker/jupyter/unsloth_labext/src/outputSelect.ts index a1641a967f..d7329f2215 100644 --- a/docker/jupyter/unsloth_labext/src/outputSelect.ts +++ b/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 = { 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; } diff --git a/docker/unsloth_colab_compat.py b/docker/unsloth_colab_compat.py index be3741fc77..9899979d6b 100644 --- a/docker/unsloth_colab_compat.py +++ b/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: diff --git a/tests/studio/test_branding_guard.py b/tests/studio/test_branding_guard.py index 2956220ee2..46f94ebac5 100644 --- a/tests/studio/test_branding_guard.py +++ b/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") diff --git a/tests/validate_studio_features.py b/tests/validate_studio_features.py index f4117d3010..85e3310001 100644 --- a/tests/validate_studio_features.py +++ b/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") # --------------------------------------------------------------------------