From af93868760894958d3325775f3e2547d9bfa9c81 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Mon, 6 Jul 2026 23:38:29 -0700 Subject: [PATCH 01/19] Fix repeated base model downloads across checkpoint exports (#6896) * Fix repeated base model downloads across checkpoint exports (#6890) Pre-warm the HF hub cache with the 16bit base weights before merge_and_overwrite_lora runs. The merge fetches shards with hf_hub_download(local_dir=...), which never populates the hub cache, so temporary merge directories (GGUF checkpoint exports) forced a full re-download of the base model for every checkpoint. The first export now downloads once into the cache and later exports copy from it. Skips itself when already cached, offline, on Kaggle/Colab, for local or nf4/fp4 bases, non-downloading save methods, or low disk. Opt out with UNSLOTH_PREWARM_HUB_CACHE=0. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Show MB for small base models in the pre-warm download message * Harden pre-warm: getattr for model config, abspath for relative HF_HUB_CACHE - Read config._name_or_path via getattr so a model without a config skips cleanly instead of taking the outer error path. - abspath the cache probe so a relative HF_HUB_CACHE walks up to a real root rather than "", which would zero the free-space check and skip pre-warm. Both from PR review; each covered by a test that fails without the fix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pre-warm the live-env HF cache so runtime redirects still hit (#6890) Resolve the hub cache the same way the merge does (unsloth_zoo _active_caches, live env) instead of huggingface_hub's import-time-frozen constants.HF_HUB_CACHE, and pass it as cache_dir to the cached probe, disk check and snapshot_download. Without this, a runtime HF_HOME/HF_HUB_CACHE redirect (unsloth_zoo redirect_hf_cache_if_readonly on a read-only default cache, or Studio) makes the pre-warm populate a different directory than the one the merge reads, so the cache-copy fast path misses and the base re-downloads on every export anyway. Adds 3 regression tests covering the cache_dir threading and the redirect case. * Apply ruff-format kwarg spacing to the pre-warm cache-dir changes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pre-warm the 16bit sibling for FP8 bases so their merged_16bit exports reuse the cache too For a merged_16bit export of an FP8 base with an existing 16bit sibling, the merge swaps to the sibling and downloads that (unsloth_zoo _resolve_fp8_16bit_sibling), so pre-warming the FP8 repo missed the cache and re-downloaded the sibling every export. Mirror the swap and pre-warm the sibling. No sibling still caches the FP8 repo for the in-place dequant path. Adds 2 regression tests. * Filter pre-warm shards through the safetensors index like the merge does Repos that ship a leftover shard set the index does not reference (e.g. granite-3.2) made the disk gate over-count and snapshot_download fetch shards the merge never reads. Mirror the merge: on the download path, keep only index-referenced shards. Runs after the already-cached check so the cached fast path stays network-free. Adds 2 tests. * Tighten pre-warm comments --------- Co-authored-by: Unsloth Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../test_prewarm_base_model_hub_cache.py | 459 ++++++++++++++++++ unsloth/save.py | 177 +++++++ 2 files changed, 636 insertions(+) create mode 100644 tests/saving/test_prewarm_base_model_hub_cache.py diff --git a/tests/saving/test_prewarm_base_model_hub_cache.py b/tests/saving/test_prewarm_base_model_hub_cache.py new file mode 100644 index 0000000000..cbb52863ba --- /dev/null +++ b/tests/saving/test_prewarm_base_model_hub_cache.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for #6890: repeated base-model downloads across checkpoint exports. + +merge_and_overwrite_lora downloads missing 16-bit shards with hf_hub_download(local_dir), +which never populates the persistent HF hub cache; a temporary merge directory (Studio +GGUF exports delete it) means every checkpoint export re-downloads the full base model. +_prewarm_base_model_hub_cache snapshot-downloads the base into the hub cache first so +the zoo's cache-copy fast path is hit on later exports. + +unsloth.save cannot be imported on GPU-less hosts, so these tests extract the helper's +source via ast and exec it against fakes, mirroring the other GPU-free tests. +""" + +from __future__ import annotations + +import ast +import json +import os +import types +from pathlib import Path + +import pytest + + +_SAVE_PY = Path(__file__).resolve().parent.parent.parent / "unsloth" / "save.py" +_SOURCE = _SAVE_PY.read_text(encoding = "utf-8") + + +def _extract_function(name: str) -> str: + tree = ast.parse(_SOURCE) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return ast.get_source_segment(_SOURCE, node) + raise AssertionError(f"{name} not found in unsloth/save.py") + + +class _FakePeftModel: + def __init__(self, name_or_path = "unsloth/gemma-4-31b-it-bnb-4bit"): + self.config = types.SimpleNamespace(_name_or_path = name_or_path) + + +class _Recorder: + """Callable that records calls and returns/raises per configuration.""" + + def __init__( + self, + result = None, + exc = None, + results_fn = None, + ): + self.calls = [] + self.result = result + self.exc = exc + self.results_fn = results_fn + + def __call__(self, *args, **kwargs): + self.calls.append((args, kwargs)) + if self.exc is not None: + raise self.exc + if self.results_fn is not None: + return self.results_fn(*args, **kwargs) + return self.result + + +def _build_env( + monkeypatch, + tmp_path, + shards = None, + cached = False, + free_bytes = 10**15, + base_source = None, + kaggle = False, + colab = False, + hub_cache = None, + live_hub_cache = "__same__", + fp8_sibling = None, + sibling_source = None, + index_weight_map = None, +): + """Exec the extracted helper with stubbed collaborators; returns (fn, stubs).""" + shards = ( + shards + if shards is not None + else [ + ("model-00001-of-00002.safetensors", 30 * 1024**3), + ("model-00002-of-00002.safetensors", 29 * 1024**3), + ] + ) + if base_source is None: + base_source = ("unsloth/gemma-4-31b-it", False, None, False, None) + + class _FS: + def __init__(self, token = None): + pass + + def ls( + self, + repo, + detail = True, + ): + return [{"name": f"{repo}/{n}", "size": s} for n, s in shards] + + class _LocalMiss(Exception): + pass + + hf_hub_download = _Recorder(result = str(tmp_path / "cached")) + if not cached: + hf_hub_download.exc = _LocalMiss("not cached") + # Serve model.safetensors.index.json (the merge's shard filter) while shard cache + # probes still miss, so the index-filter path can be exercised without a network. + if index_weight_map is not None: + _idx_path = tmp_path / "model.safetensors.index.json" + _idx_path.write_text(json.dumps({"weight_map": index_weight_map})) + + def _hub_dl( + repo_id = None, + filename = None, + **kw, + ): + if filename == "model.safetensors.index.json": + return str(_idx_path) + raise _LocalMiss("not cached") + + hf_hub_download.exc = None + hf_hub_download.results_fn = _hub_dl + snapshot_download = _Recorder() + determine_base_model_source = _Recorder(result = base_source) + # For the FP8 -> 16bit sibling swap: return the sibling's (16bit) source when the + # helper re-resolves the sibling, else the original base source. + if fp8_sibling is not None: + _sib_src = sibling_source or (fp8_sibling, False, None, False, None) + determine_base_model_source.results_fn = ( + lambda name, token = None: _sib_src if name == fp8_sibling else base_source + ) + resolve_fp8_16bit_sibling = _Recorder(result = fp8_sibling) + + cache_dir = tmp_path / "hub_cache" + cache_dir.mkdir(exist_ok = True) + + hf_module = types.SimpleNamespace( + HfFileSystem = _FS, + hf_hub_download = hf_hub_download, + snapshot_download = snapshot_download, + constants = types.SimpleNamespace( + HF_HUB_CACHE = hub_cache if hub_cache is not None else str(cache_dir) + ), + ) + zoo_module = types.SimpleNamespace( + determine_base_model_source = determine_base_model_source, + _resolve_fp8_16bit_sibling = resolve_fp8_16bit_sibling, + ) + # Stub the live-env cache resolver the pre-warm uses (matches what the merge reads). + _live = ( + (hub_cache if hub_cache is not None else str(cache_dir)) + if live_hub_cache == "__same__" + else live_hub_cache + ) + hf_cache_module = types.SimpleNamespace(_active_caches = lambda: (None, _live, None)) + monkeypatch.setitem(__import__("sys").modules, "huggingface_hub", hf_module) + monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.saving_utils", zoo_module) + monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.hf_cache", hf_cache_module) + + fake_shutil = types.SimpleNamespace( + disk_usage = lambda path: types.SimpleNamespace(free = free_bytes) + ) + + prints = [] + namespace = { + "os": os, + "shutil": fake_shutil, + "PeftModel": _FakePeftModel, + "get_model_name": lambda name, load_in_4bit: name.removesuffix("-bnb-4bit"), + "IS_KAGGLE_ENVIRONMENT": kaggle, + "IS_COLAB_ENVIRONMENT": colab, + "print": lambda *a, **k: prints.append(" ".join(str(x) for x in a)), + } + exec( + compile(_extract_function("_prewarm_base_model_hub_cache"), str(_SAVE_PY), "exec"), + namespace, + ) + stubs = types.SimpleNamespace( + snapshot_download = snapshot_download, + hf_hub_download = hf_hub_download, + determine_base_model_source = determine_base_model_source, + resolve_fp8_16bit_sibling = resolve_fp8_16bit_sibling, + prints = prints, + ) + return namespace["_prewarm_base_model_hub_cache"], stubs + + +def test_downloads_base_into_hub_cache(monkeypatch, tmp_path): + monkeypatch.delenv("UNSLOTH_PREWARM_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = "merged_16bit", token = "tok") + assert len(stubs.snapshot_download.calls) == 1 + _, kwargs = stubs.snapshot_download.calls[0] + assert kwargs["repo_id"] == "unsloth/gemma-4-31b-it" + # No local_dir: the whole point is populating the persistent cache. + assert "local_dir" not in kwargs + assert "model-00001-of-00002.safetensors" in kwargs["allow_patterns"] + assert "model.safetensors.index.json" in kwargs["allow_patterns"] + + +def test_skips_download_when_already_cached(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path, cached = True) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + # The cached check must not hit the network. + assert all(kwargs.get("local_files_only") for _, kwargs in stubs.hf_hub_download.calls) + + +def test_skips_when_disk_too_small_for_cache_copy(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path, free_bytes = 60 * 1024**3) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +@pytest.mark.parametrize("env_value", ["0", "false", "NO", "off"]) +def test_env_opt_out(monkeypatch, tmp_path, env_value): + monkeypatch.setenv("UNSLOTH_PREWARM_HUB_CACHE", env_value) + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +@pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) +def test_offline_skips(monkeypatch, tmp_path, var): + monkeypatch.setenv(var, "1") + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +def test_kaggle_and_colab_skip(monkeypatch, tmp_path): + for flag in ("kaggle", "colab"): + fn, stubs = _build_env(monkeypatch, tmp_path, **{flag: True}) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [], f"{flag} must skip pre-warm" + + +@pytest.mark.parametrize("save_method", ["merged_4bit", "forced_merged_4bit", "lora"]) +def test_non_downloading_save_methods_skip(monkeypatch, tmp_path, save_method): + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = save_method) + assert stubs.snapshot_download.calls == [] + + +def test_local_base_model_skips(monkeypatch, tmp_path): + local_dir = tmp_path / "local_base" + local_dir.mkdir() + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(name_or_path = str(local_dir)), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +def test_quantized_base_skips(monkeypatch, tmp_path): + fn, stubs = _build_env( + monkeypatch, + tmp_path, + base_source = ("unsloth/gemma-4-31b-it-bnb-4bit", False, None, True, "nf4"), + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +def test_non_peft_model_skips(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(object(), save_method = "merged_16bit") + assert stubs.determine_base_model_source.calls == [] + assert stubs.snapshot_download.calls == [] + + +def test_consolidated_shard_excluded_when_proper_shards_exist(monkeypatch, tmp_path): + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [ + ("consolidated.safetensors", 14 * 1024**3), + ("model-00001-of-00001.safetensors", 14 * 1024**3), + ], + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + _, kwargs = stubs.snapshot_download.calls[0] + assert "consolidated.safetensors" not in kwargs["allow_patterns"] + assert "model-00001-of-00001.safetensors" in kwargs["allow_patterns"] + + +def test_consolidated_only_repo_is_kept(monkeypatch, tmp_path): + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [("consolidated.safetensors", 14 * 1024**3)], + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + _, kwargs = stubs.snapshot_download.calls[0] + assert "consolidated.safetensors" in kwargs["allow_patterns"] + + +def test_listing_failure_is_swallowed(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path) + stubs.determine_base_model_source.exc = RuntimeError("HF is down") + fn(_FakePeftModel(), save_method = "merged_16bit") # must not raise + assert stubs.snapshot_download.calls == [] + + +def test_gpt_oss_bf16_mxfp4_swap_skips(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path) + fn( + _FakePeftModel(name_or_path = "unsloth/gpt-oss-20b-BF16"), + save_method = "mxfp4", + ) + assert stubs.snapshot_download.calls == [] + + +def test_missing_config_skips_cleanly(monkeypatch, tmp_path): + # A model whose config is None must skip silently, not fall into the outer + # exception handler that prints a misleading "Could not pre-cache" warning. + fn, stubs = _build_env(monkeypatch, tmp_path) + model = _FakePeftModel() # a PeftModel instance so the isinstance guard passes + model.config = None + fn(model, save_method = "merged_16bit") + assert stubs.determine_base_model_source.calls == [] + assert stubs.snapshot_download.calls == [] + assert not any( + "Could not pre-cache" in p for p in stubs.prints + ), "missing config took the error path instead of a clean skip" + + +def test_relative_hub_cache_does_not_falsely_skip(monkeypatch, tmp_path): + # A relative HF_HUB_CACHE whose leaf does not exist yet must still resolve to a real + # root for the disk probe; without abspath the walk-up hits "" and pre-warm is skipped. + monkeypatch.chdir(tmp_path) + fn, stubs = _build_env(monkeypatch, tmp_path, hub_cache = "relcache/hub") + fn(_FakePeftModel(), save_method = "merged_16bit") + assert len(stubs.snapshot_download.calls) == 1, "relative cache path falsely skipped pre-warm" + + +def test_generic_save_calls_prewarm_before_merge(): + """unsloth_generic_save must pre-warm the cache before merge_and_overwrite_lora.""" + tree = ast.parse(_SOURCE) + fn = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "unsloth_generic_save" + ) + body_src = ast.get_source_segment(_SOURCE, fn) + prewarm_pos = body_src.find("_prewarm_base_model_hub_cache(") + merge_pos = body_src.find("merge_and_overwrite_lora(") + assert prewarm_pos != -1, "unsloth_generic_save no longer pre-warms the hub cache" + assert merge_pos != -1 + assert prewarm_pos < merge_pos, "pre-warm must run before the merge downloads shards" + + +def test_prewarm_downloads_into_live_env_cache(monkeypatch, tmp_path): + # Download must target the live-env cache (what the merge reads), via cache_dir. + fn, stubs = _build_env(monkeypatch, tmp_path, live_hub_cache = "/mnt/persistent/hf/hub") + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/mnt/persistent/hf/hub" + + +def test_prewarm_survives_runtime_cache_redirect(monkeypatch, tmp_path): + # Frozen constants (stale dir) vs the merge's runtime-redirected dir: the pre-warm + # must follow the redirect, else the cache-copy fast path misses and #6890 is unfixed. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + hub_cache = "/read-only/original/hub", + live_hub_cache = "/writable/redirect/hub", + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/writable/redirect/hub" + + +def test_cached_probe_uses_live_env_cache(monkeypatch, tmp_path): + # The already-cached fast path must probe the live-env cache dir too. + fn, stubs = _build_env( + monkeypatch, tmp_path, cached = True, live_hub_cache = "/mnt/persistent/hf/hub" + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.hf_hub_download.calls, "cached probe did not run" + assert all( + kw.get("cache_dir") == "/mnt/persistent/hf/hub" for _, kw in stubs.hf_hub_download.calls + ) + + +def test_fp8_base_prewarms_16bit_sibling_not_fp8_repo(monkeypatch, tmp_path): + # A merged_16bit export of an FP8 base with a 16bit sibling merges onto the sibling, + # so the pre-warm must cache the sibling (what the merge downloads), not the FP8 repo. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + base_source = ("unsloth/Model-FP8", False, None, True, "fp8"), + fp8_sibling = "unsloth/Model", + ) + fn(_FakePeftModel(name_or_path = "unsloth/Model-FP8"), save_method = "merged_16bit") + assert stubs.resolve_fp8_16bit_sibling.calls, "sibling resolver was not consulted" + assert len(stubs.snapshot_download.calls) == 1 + assert stubs.snapshot_download.calls[0][1]["repo_id"] == "unsloth/Model" + + +def test_fp8_base_without_sibling_still_prewarms_fp8_repo(monkeypatch, tmp_path): + # No sibling: the merge dequants the FP8 base in place, so caching the FP8 repo helps. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + base_source = ("unsloth/Model-FP8", False, None, True, "fp8"), + fp8_sibling = None, + ) + fn(_FakePeftModel(name_or_path = "unsloth/Model-FP8"), save_method = "merged_16bit") + assert len(stubs.snapshot_download.calls) == 1 + assert stubs.snapshot_download.calls[0][1]["repo_id"] == "unsloth/Model-FP8" + + +def test_prewarm_filters_shards_through_index(monkeypatch, tmp_path): + # A repo with a leftover shard not referenced by the index: the merge keeps only the + # indexed shards, so the pre-warm must too (else the disk gate over-counts and + # snapshot_download fetches the unused leftover). + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [ + ("model-00001-of-00002.safetensors", 10 * 1024**3), + ("model-00002-of-00002.safetensors", 10 * 1024**3), + ("leftover-00001-of-00001.safetensors", 10 * 1024**3), + ], + index_weight_map = { + "a.weight": "model-00001-of-00002.safetensors", + "b.weight": "model-00002-of-00002.safetensors", + }, + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + allow = stubs.snapshot_download.calls[0][1]["allow_patterns"] + assert "leftover-00001-of-00001.safetensors" not in allow + assert "model-00001-of-00002.safetensors" in allow + assert "model-00002-of-00002.safetensors" in allow + + +def test_prewarm_keeps_all_shards_when_index_matches(monkeypatch, tmp_path): + # No leftover: every listed shard is indexed, so none are dropped. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [ + ("model-00001-of-00002.safetensors", 10 * 1024**3), + ("model-00002-of-00002.safetensors", 10 * 1024**3), + ], + index_weight_map = { + "a.weight": "model-00001-of-00002.safetensors", + "b.weight": "model-00002-of-00002.safetensors", + }, + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + allow = stubs.snapshot_download.calls[0][1]["allow_patterns"] + assert "model-00001-of-00002.safetensors" in allow + assert "model-00002-of-00002.safetensors" in allow diff --git a/unsloth/save.py b/unsloth/save.py index 020c63a9e2..0b408a1889 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -3693,6 +3693,182 @@ from unsloth_zoo.llama_cpp import ( ) +def _prewarm_base_model_hub_cache( + model, + save_method = "merged_16bit", + token = None, +): + """Download the 16-bit base weights into the persistent HF hub cache before the merge. + + merge_and_overwrite_lora fetches missing shards with hf_hub_download(local_dir = ...), + which never populates the hub cache. When the merge directory is temporary (GGUF + checkpoint exports delete it after conversion), every export re-downloads the full + base model (#6890). Pre-warming the cache makes the first export download once and + later exports copy from the cache. Best-effort: any failure or skip falls back to + the streaming download. Disable with UNSLOTH_PREWARM_HUB_CACHE=0. + """ + _false = ("0", "false", "no", "off") + if os.environ.get("UNSLOTH_PREWARM_HUB_CACHE", "1").strip().lower() in _false: + return + if IS_KAGGLE_ENVIRONMENT or IS_COLAB_ENVIRONMENT: + return + _true = ("1", "true", "yes", "on") + if ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _true + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _true + ): + return + # Only the 16bit / mxfp4 merges download the base model; merged_4bit and lora do not. + if save_method not in ("merged_16bit", "mxfp4"): + return + if not isinstance(model, PeftModel): + return + + try: + # getattr so a model without a config / _name_or_path skips instead of raising. + name_or_path = getattr(getattr(model, "config", None), "_name_or_path", None) + if not name_or_path: + return + try: + model_name = get_model_name(name_or_path, load_in_4bit = False) + except Exception: + model_name = name_or_path + if not model_name or os.path.isdir(model_name): + return # local checkpoints are copied, never downloaded + + # The merge may swap a gpt-oss "-BF16" repo for its MXFP4 variant, so skip it. + if save_method == "mxfp4" and model_name.endswith("-BF16"): + return + + from unsloth_zoo.saving_utils import determine_base_model_source + + model_name, is_local_path, _, base_is_quantized, quant_type = determine_base_model_source( + model_name, token + ) + if not model_name or is_local_path: + return + # Mirror the merge: an FP8 base with a 16bit sibling merges onto the sibling, so + # pre-warm the sibling (what the merge downloads), not the FP8 repo (#6890). + if base_is_quantized and quant_type == "fp8" and save_method == "merged_16bit": + try: + from unsloth_zoo.saving_utils import _resolve_fp8_16bit_sibling + sibling = _resolve_fp8_16bit_sibling(model_name, token) + except Exception: + sibling = None + if sibling: + model_name, is_local_path, _, base_is_quantized, quant_type = ( + determine_base_model_source(sibling, token) + ) + if not model_name or is_local_path: + return + if base_is_quantized and quant_type in ("nf4", "fp4"): + return # the 16bit merge refuses these bases; nothing worth caching + + from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download + + # Resolve the cache from the live env like the merge, not huggingface_hub's frozen + # constants: a runtime cache redirect (read-only default, Studio) would else miss (#6890). + try: + from unsloth_zoo.hf_cache import _active_caches + _hub_cache = _active_caches()[1] + hub_cache_dir = str(_hub_cache) if _hub_cache is not None else None + except Exception: + hub_cache_dir = None + + # Mirror the zoo's shard listing (drop consolidated.safetensors when proper + # shards coexist) so the cached set is a superset of what the merge looks up. + shard_names = [] + total_size_in_bytes = 0 + for x in HfFileSystem(token = token).ls(model_name, detail = True): + if x["name"].endswith(".safetensors"): + shard_names.append((os.path.split(x["name"])[-1], int(x.get("size") or 0))) + if any(name != "consolidated.safetensors" for name, _ in shard_names): + shard_names = [x for x in shard_names if x[0] != "consolidated.safetensors"] + if not shard_names: + return + + try: + for filename, _ in shard_names: + hf_hub_download( + repo_id = model_name, + filename = filename, + cache_dir = hub_cache_dir, + local_files_only = True, + token = token, + ) + return # already fully cached + except Exception: + pass + + # Mirror the merge's index filter (download path only): some repos ship leftover shards + # the index omits; keep only indexed ones, else the disk gate over-counts and we fetch + # unused shards. + if len(shard_names) > 1: + try: + import json as _json + + _idx = hf_hub_download( + repo_id = model_name, + filename = "model.safetensors.index.json", + cache_dir = hub_cache_dir, + token = token, + ) + with open(_idx, encoding = "utf-8") as _f: + _indexed = { + os.path.split(v)[-1] for v in _json.load(_f).get("weight_map", {}).values() + } + if _indexed and not {n for n, _ in shard_names}.issubset(_indexed): + _kept = [x for x in shard_names if x[0] in _indexed] + if _kept: + shard_names = _kept + except Exception: + pass + total_size_in_bytes = sum(size for _, size in shard_names) + + # The cache copy is extra disk on top of the merge working copy; need room for both. + from huggingface_hub import constants as _hf_constants + + # abspath so a relative HF_HUB_CACHE walks up to an existing root, not "". + cache_probe = os.path.abspath( + os.path.expanduser(str(hub_cache_dir or _hf_constants.HF_HUB_CACHE)) + ) + while cache_probe and not os.path.exists(cache_probe): + parent = os.path.dirname(cache_probe) + if parent == cache_probe: + break + cache_probe = parent + free_space = shutil.disk_usage(cache_probe).free if os.path.exists(cache_probe) else 0 + if free_space < 2 * total_size_in_bytes: + print( + f"Unsloth: Not enough free disk to keep `{model_name}` in the Hugging Face " + f"cache (need ~{round(2 * total_size_in_bytes / 1024**3, 1)}GB free, have " + f"{round(free_space / 1024**3, 1)}GB). Downloading straight to the merge " + f"directory instead; the next export will re-download it." + ) + return + + if total_size_in_bytes >= 0.1 * 1024**3: + size_str = f"{round(total_size_in_bytes / 1024**3, 1)}GB" + else: + size_str = f"{max(1, round(total_size_in_bytes / 1024**2))}MB" + print( + f"Unsloth: Downloading `{model_name}` into the Hugging Face cache so future " + f"exports skip the {size_str} download..." + ) + snapshot_download( + repo_id = model_name, + allow_patterns = [name for name, _ in shard_names] + + ["model.safetensors.index.json", "tokenizer.model"], + cache_dir = hub_cache_dir, + token = token, + ) + except Exception as e: + print( + f"Unsloth: Could not pre-cache the base model weights ({e}). " + f"Falling back to downloading into the merge directory." + ) + + @torch.inference_mode def save_to_gguf_generic( model, @@ -3888,6 +4064,7 @@ def unsloth_generic_save( print(f"Unsloth: Model saved successfully to '{save_directory}'") else: + _prewarm_base_model_hub_cache(model, save_method = save_method, token = token) merge_and_overwrite_lora( get_model_name, model = model, From 08226c2475e81a2c46336787f60d20a47e47ccea Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 23:48:23 -0700 Subject: [PATCH 02/19] Studio: fix torch CUDA undefined-symbol errors from a conflicting LD_LIBRARY_PATH (#6905) * Studio: re-exec to prepend torch's bundled CUDA libs to LD_LIBRARY_PATH On Linux the dynamic linker reads LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a pre-existing LD_LIBRARY_PATH pointing at a system CUDA (conda, a Docker base image, /usr/local/cuda-*/lib64) shadows torch's bundled nvidia/*/lib libraries and causes undefined-symbol errors when the Studio backend imports torch. Detect torch's lib dirs without importing torch, prepend them to LD_LIBRARY_PATH, and re-exec once (LD_LIBRARY_PATH is only read at process start). Linux-only, sentinel-guarded against re-exec loops, and called only from run.py's __main__ so library/embedder imports (e.g. Colab's `from run import run_server`) are never re-exec'd. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/run.py | 79 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/studio/backend/run.py b/studio/backend/run.py index ccd0113972..2cc6c4a93e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -12,6 +12,79 @@ import time from pathlib import Path from typing import Optional + +def _fix_torch_cuda_ld_path(): + """Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH. + + PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, ...) in + ``site-packages/nvidia/*/lib``. On Linux the dynamic linker reads + LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a + pre-existing LD_LIBRARY_PATH pointing at a different system CUDA (e.g. + /usr/local/cuda-13/lib64 from conda or a Docker base image) shadows torch's + libs and triggers "undefined symbol" errors when torch is imported. Detect + torch's lib dirs (without importing torch) and prepend them. Returns True if + LD_LIBRARY_PATH was changed. + """ + if sys.platform != "linux": + return False + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + if not ld_path: + return False + try: + import importlib.util + + spec = importlib.util.find_spec("torch") + if not spec or not spec.origin: + return False + torch_dir = os.path.dirname(spec.origin) + site_pkgs = os.path.dirname(torch_dir) + nvidia_dir = os.path.join(site_pkgs, "nvidia") + + lib_dirs = [] + torch_lib = os.path.join(torch_dir, "lib") + if os.path.isdir(torch_lib): + lib_dirs.append(torch_lib) + if os.path.isdir(nvidia_dir): + for sub in sorted(os.listdir(nvidia_dir)): + lib = os.path.join(nvidia_dir, sub, "lib") + if os.path.isdir(lib): + lib_dirs.append(lib) + if not lib_dirs: + return False + + existing = ld_path.split(":") + if existing[: len(lib_dirs)] == lib_dirs: + return False # already at the front, nothing to do + + torch_set = set(lib_dirs) + cleaned = [p for p in existing if p not in torch_set] + os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned) + return True + except Exception: + return False + + +_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED" + + +def _maybe_reexec_for_cuda_ld_path(): + """Re-exec once so the dynamic linker sees the corrected LD_LIBRARY_PATH. + + LD_LIBRARY_PATH is read at process start, so editing os.environ in-process + cannot fix the running interpreter; a single re-exec is required. Call only + from a true entry point (the ``if __name__ == "__main__"`` block), never at + import time, because os.execv replaces the whole process (an embedder such + as Colab that does ``from run import run_server`` must not be re-exec'd). + """ + if _LD_FIXED_SENTINEL in os.environ: + return + if not _fix_torch_cuda_ld_path(): + return + os.environ[_LD_FIXED_SENTINEL] = "1" + argv = getattr(sys, "orig_argv", None) or [sys.executable, *sys.argv] + os.execv(sys.executable, argv) + + # Suppress C-level dependency warnings globally (e.g. SwigPyPacked). os.environ["PYTHONWARNINGS"] = "ignore" @@ -1457,6 +1530,12 @@ def _build_arg_parser(): # For direct execution (also invoked by CLI via os.execvp / subprocess). if __name__ == "__main__": + # Correct a conflicting system CUDA on LD_LIBRARY_PATH before torch is + # imported (below, via run_server). Re-execs once on Linux so the dynamic + # linker uses torch's bundled CUDA libs; no-op on other platforms, when + # LD_LIBRARY_PATH is unset or already correct, or after the single re-exec. + _maybe_reexec_for_cuda_ld_path() + import signal import traceback From 69f8e0b228200ea8e748b5c3a118eef71f441168 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 00:06:48 -0700 Subject: [PATCH 03/19] Clear stale yolo approval state on no-launch reruns (#6868) * Clear stale yolo approval state on no-launch reruns The no-launch session config dir is deliberately reused across runs, but the config writers only ever added the --yolo auto-approval settings and never removed them. After one --yolo --no-launch run, every later run without --yolo kept OpenClaw's tools.exec security=full/ask=off policy plus exec-approvals.json, and OpenCode's permission allow block, so tool execution stayed silently pre-approved. Non-yolo runs now reset that state: OpenClaw drops the exec policy keys and the yolo defaults in exec-approvals.json (approvals OpenClaw itself recorded are kept; the file is removed when only the yolo payload is left), and OpenCode drops the permission block. Launch mode is untouched since it already uses an ephemeral temp dir. * Strip only yolo-written values on non-yolo cleanup Match each field against the exact value the yolo path writes before removing it, so a stricter exec policy, approvals defaults set by the user or the OpenClaw UI, and deny/ask OpenCode permission entries all survive a plain no-launch rerun. An unparseable exec-approvals.json is left in place, matching how an unparseable config is handled. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Write a prompting policy on non-yolo instead of deleting to a permissive default OpenClaw and OpenCode both treat an omitted policy as permissive: OpenClaw's effective exec policy for an unset tools.exec is security=full/ask=off on the gateway host, and OpenCode defaults an unset permission to allow. So clearing the yolo values on a non-yolo run did not restore prompting, it fell back to those permissive defaults and left tool execution auto-approved. A non-yolo run now writes an explicit prompting policy: OpenClaw gets security=allowlist/ask=on-miss (verified to prompt even with the approvals file removed, since the stricter of config and approvals wins), and OpenCode gets edit/bash/webfetch=ask. Only a permissive/yolo value is tightened; a stricter deny (or an ask the user set) is preserved, and the yolo approvals defaults are still stripped. The file-edit CI path opts opencode/openclaw into --yolo, since those agents now prompt by default and the headless test needs auto-approval. * Respect existing exec mode, sandbox/node host, and global permission rules on non-yolo reset The non-yolo reset for openclaw/opencode assumed an omitted policy was the permissive yolo default and rewrote it, which corrupted or weakened stricter setups it should have preserved: - OpenClaw tools.exec.mode is the normalized policy knob and cannot be combined with explicit security/ask (OpenClaw rejects the whole config), so writing security+ask alongside a mode:deny/ask policy both broke the config and relaxed it. Leave a mode-based policy untouched. - host=sandbox defaults to security=deny and host=node routes to a paired node; neither is written by --yolo (which only writes host=gateway). Treating the missing security as full and popping host broadened those into gateway/auto exec. Only rewrite a gateway-routed permissive policy, and never pop a non-gateway host. - OpenCode permission can be a string ("deny") or a {"*": ...} catch-all. The old code dropped a string form and overrode a catch-all by writing per-tool ask, weakening a stricter user rule. Now a string is left in place, a catch-all governs absent tools, and only an effective allow is tightened. - The non-yolo ask policy only lived in OPENCODE_CONFIG, which loads below project opencode.json, so a project config allowing edit/bash/webfetch still auto-approved. Carry the ask policy in OPENCODE_CONFIG_CONTENT (above project config) too, symmetric to how yolo carries its allow. Also harden the openclaw path against a malformed non-dict tools value. Adds tests for mode/sandbox/node hosts, string and catch-all permissions, and the inline ask policy over a project config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope non-yolo resets to the exact yolo fingerprint and preserve granular denies OpenClaw: reset only the exact host=gateway/security=full/ask=off policy --yolo writes, so an omitted or host=auto/sandbox/node policy (which can resolve to a sandbox security=deny default) is no longer broadened to allowlist/on-miss, and a deliberate tools.exec.mode is left alone (OpenClaw never migrates our security/ask write into a mode). OpenCode: carry a granular object or a deny inline verbatim so a per-tool user rule is not collapsed to a blanket ask, but floor any object that grants allow anywhere to the string ask (which fully replaces a project object) so no inline allow pattern can leak through into a silent auto-approve on a non-yolo session. * Stop overriding project config on non-yolo; require full approvals fingerprint The non-yolo OpenCode reset carried a session permission in OPENCODE_CONFIG_CONTENT, which outranks the project opencode.json we cannot read. That inline override could not correctly reflect the project: it weakened a project deny to a prompt, mishandled global string rules, leaked through a granular object's permissive default when no catch-all was present, collapsed an object with an allow (losing its deny), and missed per-agent permissions. All of these stem from forcing a value over an unknown project config. A non-yolo run now only undoes what --yolo wrote: it flips our own explicit per-tool allow back to ask in our config file and carries no permission inline, so the project's own permissions are honored as written. Clearing our persisted yolo state is the actual fix; --yolo still carries its allow inline so it works over a project config. OpenClaw approvals cleanup now strips the yolo defaults only when the full fingerprint (security=full, ask=off, askFallback=full) is present, so a mixed user policy that merely shares askFallback=full (whose omitted default is deny) is kept intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/scripts/agent-guides-drive.sh | 12 +- unsloth_cli/commands/start.py | 102 +++++++-- unsloth_cli/tests/test_start.py | 293 +++++++++++++++++++++++++- 3 files changed, 389 insertions(+), 18 deletions(-) diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 9b85b20177..d430d2c172 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -154,7 +154,12 @@ raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) # writers as a side effect (it writes each agent's relocated session config). parse_connect() { local raw="$LOGS_DIR/connect-${AGENT}.txt" - if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + # CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their + # config (which now prompts by default), so the file-edit test opts into auto-approval + # here, the same intent as claude/codex's per-call bypass flags. + local yolo=() + [ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo) + if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then cat_redacted "$raw" guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" fi @@ -394,7 +399,10 @@ case "$MODE" in T2='Run hello.py with python and show me the exact output.' # The start.py recipe writers + crosscheck must see the repo; run them - # from the repo root BEFORE cd-ing into the scratch work dir. + # from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw + # gate tool approval through their config (prompting by default), so file-edit + # opts them into auto-approval to run edits/commands headlessly. + case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac parse_connect crosscheck_contract # File-edit needs real tools, so we cannot zero them as in connection. diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index b188180188..1895125b11 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -983,7 +983,9 @@ def _session_config(agent: str, launch: bool): else: # Never wipe this dir: a previously printed recipe may still be running # an agent whose sessions/state live here, and every config writer - # merges idempotently into an existing home anyway. + # merges idempotently into an existing home anyway. Writers must also + # reset any state a previous run's flags left behind (--yolo especially), + # since files here outlive the invocation that wrote them. path = _agents_config_root() / agent path.mkdir(parents = True, exist_ok = True, mode = 0o700) yield path @@ -1043,6 +1045,62 @@ def write_openclaw_config( {"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}}, ) typer.echo(f"Updated {approvals}") + else: + # The no-launch config dir is reused across runs, so a previous --yolo run may + # have left auto-approval state behind. OpenClaw treats an omitted exec policy as + # security=full, ask=off on the gateway host, so deleting the keys would keep + # auto-approval on: a non-yolo run must WRITE a prompting policy. Only a + # permissive/yolo policy is replaced; a stricter one set by hand survives. + tools = config.get("tools") + exec_policy = tools.get("exec") if isinstance(tools, dict) else None + exec_policy = exec_policy if isinstance(exec_policy, dict) else {} + # Match ONLY the exact fingerprint --yolo writes (host=gateway, security=full, + # ask=off, all explicit, no mode); anything else is left untouched. host=auto or an + # omitted host resolves to security=deny under an active sandbox, so treating those + # as the permissive gateway default would broaden a fresh sandboxed config from + # deny to allowlist. host=node and host=sandbox are user-set (--yolo only writes + # gateway). tools.exec.mode is OpenClaw's normalized knob (it cannot be combined + # with security/ask, and OpenClaw never rewrites our security/ask write into it), + # so a mode is always a deliberate user policy; never clobber it. + permissive = ( + "mode" not in exec_policy + and exec_policy.get("host") == "gateway" + and exec_policy.get("security") == "full" + and exec_policy.get("ask") == "off" + ) + if permissive: + exec_policy = _subdict(_subdict(config, "tools"), "exec") + exec_policy.pop("host", None) # routing only; defaults to the gateway host + exec_policy["security"] = "allowlist" # only allowlisted commands skip approval + exec_policy["ask"] = "on-miss" # prompt on every non-allowlisted command + # Drop the yolo defaults from the host approvals file (a stricter default set by + # the user or OpenClaw is kept). With a prompting tools.exec the stricter of the + # two layers wins, so an omitted approvals default still prompts. + approvals = path.parent / "exec-approvals.json" + if approvals.exists(): + state = _read_json_object(approvals) + if state is not None: + defaults = state.get("defaults") + # Strip the defaults only when they are exactly the yolo fingerprint; a + # user-managed mixed policy that merely shares a field (e.g. askFallback=full, + # whose omitted default is deny) must be kept intact. + yolo_defaults = (("security", "full"), ("ask", "off"), ("askFallback", "full")) + is_yolo = isinstance(defaults, dict) and all( + defaults.get(k) == v for k, v in yolo_defaults + ) + if is_yolo: + for k, _ in yolo_defaults: + del defaults[k] + if not defaults: + del state["defaults"] + if set(state) <= {"version"}: + # Nothing left but our own yolo payload: remove it. + approvals.unlink() + typer.echo(f"Removed {approvals}") + else: + # Keep approvals OpenClaw itself recorded; only the yolo defaults go. + _write_private_json(approvals, state) + typer.echo(f"Updated {approvals}") if json.dumps(config, sort_keys = True) != before: _write_private_json(path, config) typer.echo(f"Updated {path}") @@ -1054,7 +1112,7 @@ def write_opencode_config( model: dict, path: Path, yolo: bool = False, -) -> None: +) -> dict: config = _read_json_object(path) if config is None: typer.echo( @@ -1062,7 +1120,7 @@ def write_opencode_config( "yourself, or move the file aside and re-run.", err = True, ) - return + return {} before = json.dumps(config, sort_keys = True) config.setdefault("$schema", "https://opencode.ai/config.json") model_entry = {"name": model["id"]} @@ -1087,13 +1145,32 @@ def write_opencode_config( compaction = _subdict(config, "compaction") compaction["auto"] = True compaction["reserved"] = max(1, window // 10) + tools = ("edit", "bash", "webfetch") if yolo: # OpenCode has no --yolo flag; auto-approve is the config `permission` block - # (singular). Allow the prompting tools so tool calls don't block on the TUI. - config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + # (singular). Allow the prompting tools so tool calls don't block on the TUI. This + # rides inline (OPENCODE_CONFIG_CONTENT) so --yolo works even over a project config. + session_permission = {t: "allow" for t in tools} + config["permission"] = dict(session_permission) + else: + # Undo only what --yolo wrote: our yolo sets an explicit per-tool "allow" for these + # three tools, so flip exactly those explicit allows back to "ask". A "deny"/"ask", + # a granular object, a string, or a "*" catch-all is the user's own rule and is left + # untouched. We do NOT carry a permission inline for a non-yolo session: since + # OPENCODE_CONFIG_CONTENT outranks the project opencode.json we cannot read, any + # value forced there would override the user's project rules (weakening a project + # deny, or auto-approving through a granular object's permissive default). Clearing + # our own persisted yolo state is the fix; the project's own permissions are honored. + session_permission: dict = {} + permission = config.get("permission") + if isinstance(permission, dict): + for tool in tools: + if permission.get(tool) == "allow": + permission[tool] = "ask" if json.dumps(config, sort_keys = True) != before: _write_private_json(path, config) typer.echo(f"Updated {path}") + return session_permission def write_hermes_config(base: str, model: dict, path: Path) -> None: @@ -1379,14 +1456,15 @@ def opencode( # OPENCODE_CONFIG is an overlay (loaded between the user's global and project # configs), so this adds the Unsloth provider/model for the session without # changing the user's default model. Key lives in the config, not the env. - write_opencode_config(base, key, entry, config_path, yolo = yolo) - # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model - # pin (and --yolo permissions) would silently lose to a repo config. Carry the - # settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project - # config; the API key stays in the private file, never in the printed env. + session_permission = write_opencode_config(base, key, entry, config_path, yolo = yolo) + # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model pin + # would silently lose to a repo config. Carry it in OPENCODE_CONFIG_CONTENT, which + # outranks project config; the API key stays in the private file, never the env. + # Only --yolo carries a permission here (its allow must win over a project config); + # a non-yolo session returns no permission, so the project's own rules are honored. inline_config: dict = {"model": f"unsloth/{entry['id']}"} - if yolo: - inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + if session_permission: + inline_config["permission"] = session_permission env = { "OPENCODE_CONFIG": str(config_path), "OPENCODE_CONFIG_CONTENT": json.dumps(inline_config), diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index a6a092a17c..affc24626e 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -446,7 +446,10 @@ def test_opencode_inline_config_beats_project_config(fake_studio): assert "sk-unsloth" not in content_line # key stays in the private file -def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): +def test_opencode_inline_config_omits_permission_without_yolo(fake_studio): + # A non-yolo session carries no permission inline. OPENCODE_CONFIG_CONTENT outranks the + # project opencode.json we cannot read, so forcing any value there would override the + # user's project rules; clearing our own config is the fix, and the inline pins the model. result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output content_line = next( @@ -455,7 +458,8 @@ def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): inline = json.loads( shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] ) - assert inline == {"model": f"unsloth/{MODEL['id']}"} + assert inline["model"] == f"unsloth/{MODEL['id']}" + assert "permission" not in inline def test_https_loopback_never_auto_serves(fake_studio, monkeypatch): @@ -1676,9 +1680,31 @@ def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + # A non-yolo run on a fresh config writes no permission block; it only flips a prior + # --yolo run's explicit allow back to ask (see the yolo-then-plain test below). assert "permission" not in config +def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path): + # The core reset: a --yolo run wrote explicit per-tool allow; a later non-yolo run + # must flip exactly those back to ask so nothing stays auto-approved. + yolo = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert yolo.exit_code == 0, yolo.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + assert json.loads(config_path.read_text())["permission"] == { + "edit": "allow", + "bash": "allow", + "webfetch": "allow", + } + plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert plain.exit_code == 0, plain.output + assert json.loads(config_path.read_text())["permission"] == { + "edit": "ask", + "bash": "ask", + "webfetch": "ask", + } + + def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) assert result.exit_code == 0, result.output @@ -1691,12 +1717,17 @@ def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"} -def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path): +def test_no_yolo_openclaw_leaves_fresh_config_untouched(fake_studio, tmp_path): + # A fresh non-yolo run only undoes state a prior --yolo wrote; with no yolo + # fingerprint present it must not synthesize an exec policy. An omitted policy can + # resolve to a sandbox default of security=deny, so writing allowlist here would + # BROADEN it. The reset is scoped to the exact yolo write, verified by the + # yolo-then-plain round trip below. result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) assert result.exit_code == 0, result.output state = tmp_path / "agents" / "openclaw" config = json.loads((state / "openclaw.json").read_text()) - assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo + assert "exec" not in config.get("tools", {}) assert not (state / "exec-approvals.json").exists() @@ -1719,6 +1750,260 @@ def test_write_openclaw_config_yolo_unit(tmp_path): } +def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp_path): + # The no-launch config dir is reused across runs, so a --yolo run persists its + # auto-approve settings; a later run without --yolo must strip them, not leave + # tool execution silently pre-approved. + yolo = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert yolo.exit_code == 0, yolo.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + assert "permission" in json.loads(config_path.read_text()) + plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert plain.exit_code == 0, plain.output + config = json.loads(config_path.read_text()) + # The yolo allow policy is replaced by a prompting one, not deleted (which would + # revert to OpenCode's permissive "allow" default). + assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} + # The session provider survives the cleanup. + assert "unsloth" in config["provider"] + + +def test_no_launch_rerun_clears_stale_openclaw_yolo_state(fake_studio, tmp_path): + yolo = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) + assert yolo.exit_code == 0, yolo.output + state = tmp_path / "agents" / "openclaw" + assert (state / "exec-approvals.json").exists() + plain = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert plain.exit_code == 0, plain.output + config = json.loads((state / "openclaw.json").read_text()) + # The yolo policy is replaced by a prompting one, not deleted (which would revert + # to OpenClaw's permissive default), and the yolo approvals file is gone. + assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"} + assert not (state / "exec-approvals.json").exists() + # The session provider survives the cleanup. + assert "unsloth" in config["models"]["providers"] + + +def test_write_openclaw_config_yolo_then_plain_unit(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + # A plain rerun replaces the yolo policy with a prompting one (deleting it would + # fall back to OpenClaw's permissive default) and removes the yolo approvals file. + assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"} + assert not (path.parent / "exec-approvals.json").exists() + + +def test_write_opencode_config_yolo_then_plain_unit(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + # A plain rerun replaces the yolo allow policy with a prompting one. + assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} + + +def test_openclaw_non_yolo_keeps_runtime_approvals(tmp_path): + # OpenClaw records its own entries in exec-approvals.json (OPENCLAW_STATE_DIR is + # this dir); the non-yolo reset drops only the yolo defaults, not those. + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + approvals = path.parent / "exec-approvals.json" + state = json.loads(approvals.read_text()) + state["agents"] = {"main": {"allowlist": ["git status"]}} + approvals.write_text(json.dumps(state)) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + remaining = json.loads(approvals.read_text()) + assert "defaults" not in remaining + assert remaining["agents"] == {"main": {"allowlist": ["git status"]}} + + +def test_openclaw_non_yolo_keeps_mixed_approval_defaults(tmp_path): + # A mixed user-managed defaults block that only shares a field with the yolo payload + # (here askFallback=full, whose omitted default is deny) is not stale yolo state, so a + # non-yolo run leaves it intact rather than stripping the shared field. + path = tmp_path / "openclaw.json" + approvals = path.parent / "exec-approvals.json" + mixed = { + "version": 1, + "defaults": {"security": "allowlist", "ask": "on-miss", "askFallback": "full"}, + } + approvals.write_text(json.dumps(mixed)) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(approvals.read_text()) == mixed + + +def test_openclaw_non_yolo_leaves_partial_policy_untouched(tmp_path): + # A policy that lacks the full yolo fingerprint (here no host and no security) is not + # our --yolo write, so a non-yolo run leaves it as-is rather than assuming ask=off + # means permissive: an omitted host/security can resolve to a sandbox deny default. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"timeout": 30, "ask": "off"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"timeout": 30, "ask": "off"} + + +def test_openclaw_non_yolo_leaves_no_permissive_values(tmp_path): + # The whole point of the reset: after a yolo run, a plain run must leave neither the + # config nor the approvals file at OpenClaw's permissive (security=full, ask=off) + # default, or exec still auto-approves. + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + exec_policy = json.loads(path.read_text())["tools"]["exec"] + assert exec_policy.get("security") != "full" + assert exec_policy.get("ask") != "off" + assert not (path.parent / "exec-approvals.json").exists() + + +def test_openclaw_non_yolo_preserves_stricter_exec_policy(tmp_path): + # A policy that doesn't carry the yolo values (for example stricter security or + # prompting turned on) was not written by --yolo and must survive a plain run. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"security": "deny", "ask": "on"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"security": "deny", "ask": "on"} + + +def test_openclaw_non_yolo_preserves_stricter_approval_defaults(tmp_path): + # exec-approvals.json defaults that don't match the yolo payload (stricter + # settings from the user or the OpenClaw UI) are kept, and the file stays. + path = tmp_path / "openclaw.json" + approvals = path.parent / "exec-approvals.json" + approvals.write_text( + json.dumps({"version": 1, "defaults": {"security": "allowlist", "ask": "on"}}) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + state = json.loads(approvals.read_text()) + assert state["defaults"] == {"security": "allowlist", "ask": "on"} + + +def test_openclaw_non_yolo_leaves_unparseable_approvals(tmp_path): + # An unreadable approvals file is left in place rather than deleted, matching + # how an unparseable config is handled. + path = tmp_path / "openclaw.json" + approvals = path.parent / "exec-approvals.json" + approvals.write_text("{not json") + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert approvals.read_text() == "{not json" + + +def test_opencode_non_yolo_flips_only_explicit_allow(tmp_path): + # Only a tool explicitly set to "allow" (what --yolo writes) is flipped to "ask". A + # deny/ask a user set is kept, and an absent tool is not added. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"permission": {"edit": "allow", "bash": "deny", "read": "ask"}})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["permission"] == {"edit": "ask", "bash": "deny", "read": "ask"} + assert session == {} # a non-yolo session carries no permission inline + + +def test_opencode_non_yolo_leaves_string_permission(tmp_path): + # A global string rule ("deny") is a user-managed catch-all; leave it untouched and + # carry no inline override. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"permission": "deny"})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(path.read_text())["permission"] == "deny" + assert session == {} + + +def test_opencode_non_yolo_leaves_catch_all_and_flips_explicit_allow(tmp_path): + # A "*" catch-all is the user's own rule, never something --yolo writes (yolo sets + # explicit per-tool allow), so it is left intact; an explicit per-tool "allow" is still + # flipped to "ask", but an absent tool inheriting the catch-all is not touched. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"permission": {"*": "allow", "bash": "allow"}})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(path.read_text())["permission"] == {"*": "allow", "bash": "ask"} + assert session == {} + + +def test_opencode_non_yolo_leaves_granular_object(tmp_path): + # A granular object value is a user rule (yolo only ever writes a plain "allow" string), + # so it is left in the file verbatim and never carried inline. + path = tmp_path / "opencode.json" + obj = {"read *": "deny", "git *": "ask"} + path.write_text(json.dumps({"permission": {"bash": dict(obj)}})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(path.read_text())["permission"]["bash"] == obj + assert session == {} + + +def test_openclaw_non_yolo_leaves_mode_policy(tmp_path): + # tools.exec.mode is OpenClaw's normalized knob and cannot be combined with explicit + # security/ask (the config is rejected), so a mode-based policy must be left as-is. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"mode": "deny"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"mode": "deny"} + + +def test_openclaw_non_yolo_preserves_sandbox_host(tmp_path): + # host=sandbox defaults to security=deny (stricter than the gateway "full" default), + # so a non-yolo run must not treat the missing security as permissive nor pop host + # (which would broaden routing to the gateway). + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"host": "sandbox"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "sandbox"} + + +def test_openclaw_non_yolo_preserves_node_host(tmp_path): + # host=node routes to a paired node and is only ever set by the user (--yolo writes + # host=gateway), so a non-yolo run must not pop it and reroute to the gateway. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"host": "node"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "node"} + + +def test_openclaw_non_yolo_preserves_auto_host_permissive(tmp_path): + # host=auto (or omitted) with security=full/ask=off is NOT the --yolo write: under an + # active sandbox, auto resolves to security=deny. --yolo only ever writes host=gateway, + # so the reset must not treat auto/None as the permissive gateway default and broaden a + # sandboxed deny to allowlist. + for exec_policy in ( + {"host": "auto", "security": "full", "ask": "off"}, + {"security": "full", "ask": "off"}, + ): + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": dict(exec_policy)}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == exec_policy + + +def test_openclaw_non_yolo_resets_only_gateway_yolo_fingerprint(tmp_path): + # The reset fires on exactly the host=gateway + security=full + ask=off write --yolo + # makes, and nothing else. + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps({"tools": {"exec": {"host": "gateway", "security": "full", "ask": "off"}}}) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"} + + +def test_openclaw_non_yolo_preserves_full_mode(tmp_path): + # OpenClaw never normalizes our security=full/ask=off yolo write into mode:"full" + # (verified against the binary: doctor --fix and config get leave security/ask as-is), + # so a mode:"full" is always a deliberate user policy, not stale yolo state; leave it. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"mode": "full"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"mode": "full"} + + def test_yolo_command_flags_unmapped_agent_is_empty(): # Config-based agents (and any typo) must yield no flag, not a KeyError. assert start._yolo_command_flags("opencode", True) == [] From 296cacb5a176a31fce91b4ef90ed1a044dfc4fa5 Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Tue, 7 Jul 2026 04:29:37 -0500 Subject: [PATCH 04/19] ROCm-on-WSL: support discrete Radeon (RDNA 3/4) in WSL, not just Strix Halo (#6915) * WSL ROCm: generalize ROCm-on-WSL bootstrap from Strix-only to any RDNA arch install_rocm_wsl_strixhalo.sh hardcoded gfx1151, so its verify step died on discrete Radeon cards even though the ROCm + librocdxg setup is arch-agnostic. Auto-detect the GPU arch from rocminfo (override via UNSLOTH_WSL_GFX), verify any GPU agent enumerates over DXG, and map the arch to AMD's per-arch wheel family for the optional smoke test (injecting librocdxg into torch/lib so torch's bundled ROCr finds the DXG bridge). Verified on gfx1200 (Radeon RX 9060 XT) in WSL2 + Ubuntu 24.04 -- torch.cuda now enumerates the GPU. * WSL ROCm: trigger the ROCm-on-WSL bootstrap for discrete Radeon GPUs too _maybe_bootstrap_rocm_wsl only fired for Strix APUs (matched via /proc/cpuinfo, which discrete cards don't appear in). Add _wsl_amd_gpu_name() -- queries the Windows host via WMI -- and broaden the trigger gate plus the 'already-usable ROCm' rocminfo check from gfx1151-only to any real GPU agent (gfxNNNN, excluding the gfx11-generic fallback ISA). The generalized bootstrap then auto-detects the arch. Enables 'curl install.sh | sh' to set up ROCm-on-WSL on discrete Radeon RX 7000/9000 in WSL2 + Ubuntu 24.04, not just Strix Halo/Point. * WSL ROCm: address review -- filter generic ISA in bootstrap, bound the host GPU query - install_rocm_wsl_strixhalo.sh: exclude the gfx11-generic fallback ISA in arch detection (grep -v generic), matching install.sh's rocminfo check, so a generic agent listed before the real one can't be picked as the arch. - install.sh: wrap the powershell.exe Win32_VideoController query in _run_bounded (10s timeout) so an unstable WSL-interop / busy host can't hang the installer. * WSL ROCm: harden arch-detect + librocdxg copy under set -eo pipefail (review) - _detected_gfx: append '|| true' so a no-GPU rocminfo (empty pipeline, non-zero under pipefail) doesn't abort the assignment before the '[ -z ]' branch prints the diagnostic + die message. - smoke-test librocdxg copy: gate on '[ -d "$_tlib" ]' instead of '[ -n ]' so a non-directory value can't make cp rename librocdxg to 'lib'. * WSL ROCm: address Codex review (gfx000, 24.04 reroute for discrete, test locator) - Exclude gfx000 (the CPU agent) from the WSL 'usable ROCm' check and the bootstrap arch-detect: match gfx[1-9] (nonzero arch), so a partial ROCm install that only reports the CPU ISA no longer short-circuits the librocdxg setup. (P2) - Reuse the Ubuntu-24.04 reroute for discrete Radeon: broaden _maybe_reroute_strixhalo_to_2404's gate with the same _wsl_amd_gpu_name (WMI) fallback, so a discrete card on 26.04 reroutes to a 24.04 distro like Strix does instead of falling to CPU. Moved _wsl_amd_gpu_name above the reroute and made it self-contained + 10s-bounded (it runs before _run_bounded is defined). (P2) - Update TestInstallShDropinPersistence to locate the gate by its unique '!/generic/' clause now that the gfx1151 literal is gone. (P1) * Condense ROCm-on-WSL comments in install.sh and bootstrap helper * Guard WSL reroute from NVIDIA hybrid hosts and fix GFX-override pipefail check * Honor CUDA_VISIBLE_DEVICES-hidden NVIDIA in the WSL reroute guard * Reuse _has_usable_nvidia_gpu in the WSL reroute guard --------- Co-authored-by: Daniel Han --- install.sh | 162 +++++++++++++--------- scripts/install_rocm_wsl_strixhalo.sh | 57 +++++--- tests/studio/install/test_rocm_support.py | 48 ++++++- 3 files changed, 182 insertions(+), 85 deletions(-) diff --git a/install.sh b/install.sh index 81c50bc899..14fbba478d 100755 --- a/install.sh +++ b/install.sh @@ -1483,6 +1483,81 @@ elif [ "$OS" = "macos" ]; then fi tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" +# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in +# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded +# to 10s. Defined here so the reroute below can use it before _run_bounded exists. +_WSL_AMD_GPU_NAME_CACHE="" +_wsl_amd_gpu_name() { + if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then + [ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1 + printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0 + fi + command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; } + _wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name" + if command -v timeout >/dev/null 2>&1; then + _wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + else + _wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + fi + if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi + _WSL_AMD_GPU_NAME_CACHE="-"; return 1 +} + +# ── Bounded command runner ── +# Runs a command under a 10s timeout when the `timeout` binary is available, +# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during +# driver init or after a reset) from hanging the installer: a timed-out probe +# exits nonzero and is treated exactly like a failed probe. No-op semantics on +# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +_run_bounded() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every +# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to +# the AMD card). Unset means all devices visible. nvidia-smi ignores this env +# var, so the probes below cannot see the distinction on their own. +_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] +} + +# ── NVIDIA usable-GPU helper ── +# Returns 0 (true) if an NVIDIA GPU is present and usable. +# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, +# which the NVIDIA driver populates on Linux regardless of nvidia-smi state +# -- handles PATH gaps, subprocess timeouts, and driver init races that +# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. +# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. +_has_usable_nvidia_gpu() { + if _cvd_hides_nvidia; then + return 1 + fi + _nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _nvsmi="/usr/bin/nvidia-smi" + fi + if [ -n "$_nvsmi" ]; then + if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + # Fallback: NVIDIA driver exposes one subdir per GPU under this path. + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 +} + # Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04) # with a 24.04 distro present, re-run the install there and stop; else fall through # to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before @@ -1493,7 +1568,15 @@ _maybe_reroute_strixhalo_to_2404() { [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 [ -e /dev/dxg ] || return 0 - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on + # this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors + # CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps. + if _has_usable_nvidia_gpu; then return 0; fi + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi # Already ROCm-on-WSL? leave a working GPU alone, whatever the version. if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then return 0 @@ -1959,61 +2042,6 @@ _has_amd_rocm_gpu() { return 1 } -# ── Bounded command runner ── -# Runs a command under a 10s timeout when the `timeout` binary is available, -# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during -# driver init or after a reset) from hanging the installer: a timed-out probe -# exits nonzero and is treated exactly like a failed probe. No-op semantics on -# hosts without `timeout` (e.g. macOS) or when the probe is healthy. -_run_bounded() { - if command -v timeout >/dev/null 2>&1; then - timeout 10 "$@" - else - "$@" - fi -} - -# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every -# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to -# the AMD card). Unset means all devices visible. nvidia-smi ignores this env -# var, so the probes below cannot see the distinction on their own. -_cvd_hides_nvidia() { - [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 - _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') - [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] -} - -# ── NVIDIA usable-GPU helper ── -# Returns 0 (true) if an NVIDIA GPU is present and usable. -# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, -# which the NVIDIA driver populates on Linux regardless of nvidia-smi state -# -- handles PATH gaps, subprocess timeouts, and driver init races that -# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. -# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches -# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. -_has_usable_nvidia_gpu() { - if _cvd_hides_nvidia; then - return 1 - fi - _nvsmi="" - if command -v nvidia-smi >/dev/null 2>&1; then - _nvsmi="nvidia-smi" - elif [ -x "/usr/bin/nvidia-smi" ]; then - _nvsmi="/usr/bin/nvidia-smi" - fi - if [ -n "$_nvsmi" ]; then - if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then - return 0 - fi - fi - # Fallback: NVIDIA driver exposes one subdir per GPU under this path. - if [ -d /proc/driver/nvidia/gpus ] && \ - [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then - return 0 - fi - return 1 -} - # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver @@ -2327,19 +2355,19 @@ _persist_rocm_wsl_dropin() { fi } +# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it. _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 # Leave any already-usable GPU completely alone (NVIDIA, or working ROCm). if _has_usable_nvidia_gpu; then return 0; fi - # "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the - # generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and - # would skip this bootstrap while the real GPU is still unusable. awk consumes - # all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail. + # Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000, + # the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so + # rocminfo isn't SIGPIPE'd like `grep -q` under pipefail. _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then # rocminfo may work only via the transient env _ensure_rocm_probe_env # just set, which dies with the installer. Persist the drop-in so login # shells (Studio, llama.cpp) inherit it -- else a reinstall over an @@ -2349,9 +2377,12 @@ _maybe_bootstrap_rocm_wsl() { fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). [ -e /dev/dxg ] || return 0 - # Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match - # the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S"). - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also + # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi command -v bash >/dev/null 2>&1 || return 0 # Fast path: already configured (librocdxg present) but launched from a @@ -2369,7 +2400,8 @@ _maybe_bootstrap_rocm_wsl() { fi echo "" - substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN" + _rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU" + substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN" substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU." substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)" diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh index 5ef9ee386a..aa560fc432 100644 --- a/scripts/install_rocm_wsl_strixhalo.sh +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -3,13 +3,14 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # # ────────────────────────────────────────────────────────────────────────────── -# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151) +# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX +# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT). # ────────────────────────────────────────────────────────────────────────────── -# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime -# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG -# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04 -# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via -# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies). +# install.sh routes the detected arch to the right ROCm wheels once a runtime exists; +# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg). +# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by +# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the +# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent. # # Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with # production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once @@ -34,10 +35,12 @@ set -euo pipefail # ── Tunables (override via env) ────────────────────────────────────────────── ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install -GFX="gfx1151" +# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200). +# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch. +GFX="${UNSLOTH_WSL_GFX:-}" LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build -# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test. -TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/" +# AMD's wheel index for the (optional) smoke test; resolved after arch detection. +TORCH_INDEX="" # Optional torch smoke test (throwaway venv). OFF by default: install.sh installs # torch itself into the real venv right after, so a duplicate download is wasteful. SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}" @@ -220,12 +223,12 @@ $SUDO ldconfig say "Persisting ROCm-on-WSL environment" _envfile="/etc/profile.d/unsloth-rocm-wsl.sh" $SUDO tee "$_envfile" >/dev/null <>> Unsloth ROCm-on-WSL (gfx1151) >>> +# >>> Unsloth ROCm-on-WSL >>> export HSA_ENABLE_DXG_DETECTION=1 export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 export PATH="${ROCM_DIR}/bin:\${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}" -# <<< Unsloth ROCm-on-WSL (gfx1151) <<< +# <<< Unsloth ROCm-on-WSL <<< EOF # also drop into ~/.bashrc for interactive shells if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then @@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}" # ── Step 5: verify the runtime enumerates the GPU ──────────────────────────── -say "Verifying rocminfo sees ${GFX}" +say "Verifying rocminfo enumerates the GPU over DXG" # Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs # rocminfo on first match, which under `set -o pipefail` turns a successful match -# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a -# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass. +# into a pipeline failure. _rocminfo_out="$(rocminfo 2>/dev/null || true)" -if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then +# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU +# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch. +_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)" +if [ -z "$_detected_gfx" ]; then printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true - die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." + die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." fi +# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under +# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt. +if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then + die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'." +fi +GFX="${GFX:-$_detected_gfx}" # Display-only summary: best-effort (|| true) so head's early pipe-close under # `set -o pipefail` can't fail the bootstrap after verification already passed. printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true note "ROCm-on-WSL runtime is live for ${GFX}." -# ── Step 6 (optional): torch smoke test from the gfx1151 index ─────────────── +# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ─────── if [ "$SMOKE_TEST" = "1" ]; then say "Smoke-testing PyTorch on ${GFX} (throwaway venv)" + # Map the detected arch to AMD's repo.amd.com wheel family index. + case "$GFX" in + gfx1200|gfx1201) _fam="gfx120X-all" ;; + gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;; + *) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index + esac + TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/" _venv="${HOME}/.unsloth/rocm-smoketest" rm -rf "$_venv"; python3 -m venv "$_venv" "$_venv/bin/pip" install --quiet --upgrade pip - # gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py + # AMD arch index is primary (torch + triton); PyPI only an extra for pure-py # deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch. "$_venv/bin/pip" install --index-url "$TORCH_INDEX" \ --extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \ die "torch install from ${TORCH_INDEX} failed." + # WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib. + _tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)" + [ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true "$_venv/bin/python" - <<'PY' import torch ok = torch.cuda.is_available() diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index c8f2053946..bfc8132683 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3426,8 +3426,9 @@ class TestInstallShDropinPersistence: def test_gate5_early_return_persists_dropin(self): """The rocminfo-already-works early return must call the persist helper before returning.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") - # The persist call must precede `return 0` at the rocminfo gfx1151 gate. - gate = source.find("Name:[[:space:]]*gfx1151") + # The persist call must precede `return 0` at the rocminfo GPU-agent gate + # (uniquely identified by the `!/generic/` clause the other probes lack). + gate = source.find("Name:[[:space:]]*gfx[1-9]/ && !/generic/") assert gate != -1 window = source[gate : gate + 900] assert "_persist_rocm_wsl_dropin" in window @@ -3441,6 +3442,49 @@ class TestInstallShDropinPersistence: assert "profile.d/unsloth-rocm-wsl.sh" in body +_STRIXHALO_WSL_PATH = PACKAGE_ROOT / "scripts" / "install_rocm_wsl_strixhalo.sh" + + +class TestWslRerouteNvidiaGuard: + """_maybe_reroute_strixhalo_to_2404 must skip the AMD reroute on hybrid AMD+NVIDIA hosts by + reusing _has_usable_nvidia_gpu (CUDA_VISIBLE_DEVICES-aware + /proc/driver/nvidia fallback), + which must be defined before the reroute's call site so it is actually available.""" + + def test_reroute_calls_nvidia_helper_before_amd_signal(self): + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + start = source.find("_maybe_reroute_strixhalo_to_2404()") + assert start != -1 + body = source[start : start + 1200] + nv = body.find("_has_usable_nvidia_gpu") + wmi = body.find("_wsl_amd_gpu_name") + assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute" + assert wmi != -1 + # The NVIDIA guard must precede the AMD/WMI signal and return early. + assert nv < wmi + assert body.find("return 0", nv) < wmi + + def test_nvidia_helper_and_deps_defined_before_reroute_callsite(self): + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + call = source.find("\n_maybe_reroute_strixhalo_to_2404 || true") + assert call != -1 + for fn in ("_run_bounded() {", "_cvd_hides_nvidia() {", "_has_usable_nvidia_gpu() {"): + idx = source.find(fn) + assert idx != -1 and idx < call, f"{fn} must be defined before the reroute call" + + +class TestStrixhaloGfxOverridePipefail: + """The UNSLOTH_WSL_GFX override check must use a consuming grep, not grep -q: under + `set -o pipefail` an early -q exit SIGPIPEs printf and misreports the arch on large output.""" + + def test_gfx_override_uses_consuming_grep(self): + source = _STRIXHALO_WSL_PATH.read_text(encoding = "utf-8") + idx = source.find('grep -E "Name:[[:space:]]*${GFX}') + assert idx != -1, "GFX override must use a consuming grep -E (not grep -q)" + line = source[idx : source.find("\n", idx)] + assert ">/dev/null" in line + assert 'grep -qE "Name:[[:space:]]*${GFX}' not in source + + class TestLlamaCppRuntimeWslOrdering: """The serve-time launcher mirrors binary_env: system HIP before the bundle dir on WSL.""" From bdb958e052eca6a47c17410558562f00481f71b8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:16:57 -0700 Subject: [PATCH 05/19] Guard RoPE scaling against the transformers v5 buffer blank; honor extended RoPE factor (#6925) * Guard RoPE scaling against the transformers v5 buffer blank; honor extended factor Add a family-agnostic guard that builds each rotary from a scaled config, blanks its non-persistent buffers (what transformers v5 does on load), runs loader._fix_rope_inv_freq, and asserts every buffer is restored to its scaled value (llama3 and longrope). This catches the whole bug class, not just the one call site, and is validated to fail on the pre-fix repair. Also make LlamaExtendedRotaryEmbedding read the llama3 factor from the config instead of hardcoding 8 (wrong for Llama-3.2, factor 32), falling back to the Llama-3.1 defaults when built without a config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pass config into extended rotary codegen; skip v5 round-trip on transformers 4.x - patch_llama_rope_scaling now builds the llama3 extended rotary with config=self.config so it reads the real factor (32 for Llama-3.2) instead of falling back to 8; the template already references self.config. - test_v5_blank_repair_roundtrip now skips when loader._NEEDS_ROPE_FIX is False, since _fix_rope_inv_freq is a no-op on transformers 4.x and cannot restore the blanked buffers there. * Raise stream deadlock-guard timeouts from 0.2s to 5.0s in passthrough tests These asyncio.wait_for guards bound test setup and cross-task event signaling that complete near-instantly on success; the 0.2s budget is a latency assertion in disguise and times out under CI scheduling load (seen on the 3.11 matrix leg while 3.10/3.12/3.13 pass the same commit). 5.0s matches the timeout used elsewhere in the suite and still fails fast on a real hang. No test relies on the guard expiring. * Extended rotary reads rope_parameters as well as rope_scaling transformers v5 stores llama3 scaling under config.rope_parameters and exposes rope_scaling only as a back-compat property. Reading that property works on 5.0-5.13 (verified: factor resolves to 32 for Llama-3.2), but a future release may drop the shim, after which the subclass path would fall back to factor 8. Read either field so the factor survives the rename. Adds test_extended_rotary_reads_rope_parameters_v5 (fails on the old single-field read: rope_parameters-only config resolves to 8, not 32). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../tests/test_openai_tool_passthrough.py | 18 +-- tests/utils/test_rope_scaling_drift.py | 138 ++++++++++++++++++ unsloth/models/_utils.py | 1 + unsloth/models/llama.py | 17 ++- 4 files changed, 160 insertions(+), 14 deletions(-) diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index ccbd78e2b1..05e017ba7f 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1900,7 +1900,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2044,7 +2044,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) gate.set() @@ -2107,7 +2107,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2190,7 +2190,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2252,7 +2252,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY @@ -2323,13 +2323,13 @@ class TestApiMonitorProviderAndCompletionStreams: monitor_id = monitor_id, ) ) - await asyncio.wait_for(entered.wait(), timeout = 0.2) + await asyncio.wait_for(entered.wait(), timeout = 5.0) assert cancel_id in inf_mod._CANCEL_REGISTRY task.cancel() with pytest.raises(asyncio.CancelledError): await task - await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + await asyncio.wait_for(cancelled.wait(), timeout = 5.0) assert cancel_id not in inf_mod._CANCEL_REGISTRY asyncio.run(_run()) @@ -2389,13 +2389,13 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY gate.set() - await asyncio.wait_for(returned.wait(), timeout = 0.2) + await asyncio.wait_for(returned.wait(), timeout = 5.0) await asyncio.sleep(0) await response._unstarted_cleanup() assert upstream_response.is_closed diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 98f7e2db62..7a738e236c 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -257,6 +257,63 @@ def test_recompute_helper_scales_on_cpu(): ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." +def test_extended_rotary_reads_config_factor(): + # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8 + # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405). + from types import SimpleNamespace + + from unsloth.models.llama import LlamaExtendedRotaryEmbedding + + rot = object.__new__(LlamaExtendedRotaryEmbedding) + rot.base = ROPE_THETA + rot.dim = HEAD_DIM + rot._unsloth_rope_config = SimpleNamespace( + rope_scaling = { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + } + ) + vanilla = _vanilla_inv_freq() + scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1) + ratio = float(vanilla[-1]) / float(scaled[-1]) + assert abs(ratio - 32.0) < 1e-3, ( + f"LlamaExtendedRotaryEmbedding ignored config factor 32 (ratio {ratio}); the " + "low-frequency band must be divided by the config factor (issue #2405)." + ) + + +def test_extended_rotary_reads_rope_parameters_v5(): + # transformers v5 stores scaling under rope_parameters (rope_scaling is a + # back-compat shim that may be removed); the factor must still be read. + from types import SimpleNamespace + + from unsloth.models.llama import LlamaExtendedRotaryEmbedding + + rot = object.__new__(LlamaExtendedRotaryEmbedding) + rot.base = ROPE_THETA + rot.dim = HEAD_DIM + rot._unsloth_rope_config = SimpleNamespace( + rope_scaling = None, + rope_parameters = { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + }, + ) + vanilla = _vanilla_inv_freq() + scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1) + ratio = float(vanilla[-1]) / float(scaled[-1]) + assert abs(ratio - 32.0) < 1e-3, ( + f"Extended rotary ignored rope_parameters factor 32 (ratio {ratio}); v5 " + "keeps the factor under rope_parameters, not rope_scaling." + ) + + def _cos_at_position(rot, position): """cos row at one position, built like _set_cos_sin_cache but CPU-only.""" inv_freq = rot.inv_freq.float().cpu() @@ -324,6 +381,87 @@ def test_extended_cache_keeps_scaling_after_growth(): ) +def _blank_nonpersistent_buffers(module): + """Mimic transformers v5 meta-load: overwrite non-persistent buffers with garbage.""" + for name, buf in list(module.named_buffers()): + leaf = module + *parents, attr = name.split(".") + for part in parents: + leaf = getattr(leaf, part) + if attr in getattr(leaf, "_non_persistent_buffers_set", set()): + setattr(leaf, attr, torch.rand_like(buf)) + + +def _build_llama3_rotary(): + from unsloth.models import llama as llama_mod + config = _make_config(LLAMA3_ROPE_SCALING) + return llama_mod.LlamaRotaryEmbedding(config = config), config + + +def _build_longrope_rotary(): + from types import SimpleNamespace + + from unsloth.models import llama as llama_mod + + short_factor, long_factor = [1.05] * 48, [1.3] * 48 + rot = llama_mod.LongRopeRotaryEmbedding( + dim = 96, + max_position_embeddings = 131072, + original_max_position_embeddings = 4096, + base = ROPE_THETA, + short_factor = short_factor, + long_factor = long_factor, + ) + config = SimpleNamespace( + rope_scaling = { + "rope_type": "longrope", + "short_factor": short_factor, + "long_factor": long_factor, + "original_max_position_embeddings": 4096, + } + ) + return rot, config + + +@requires_cuda +@pytest.mark.parametrize( + "build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"] +) +def test_v5_blank_repair_roundtrip(build): + # Build scaled -> blank non-persistent buffers (what transformers v5 does on + # load) -> run the repair -> every buffer must return to its scaled value. + # Family-agnostic: encodes no scaling math, so it guards any rotary that + # keeps scaling in a buffer (issue #2405 / PR #6907). + from unsloth.models import loader + + # The repair only runs on transformers v5 (it is what blanks the buffers); + # on v4 _fix_rope_inv_freq is a no-op, so the round-trip cannot restore. + if not loader._NEEDS_ROPE_FIX: + pytest.skip("transformers < 5 does not blank rope buffers; repair is a no-op") + + rot, config = build() + snapshot = {name: buf.detach().clone() for name, buf in rot.named_buffers()} + assert snapshot, "rotary registers no buffers; nothing to guard" + + _blank_nonpersistent_buffers(rot) + assert any( + not torch.equal(rot.get_buffer(name), snapshot[name]) for name in snapshot + ), "blanking changed no buffer; the round-trip would be vacuous" + + wrapper = torch.nn.Module() + wrapper.add_module("rotary_emb", rot) + wrapper.config = config + loader._fix_rope_inv_freq(wrapper) + + for name in snapshot: + assert torch.allclose( + rot.get_buffer(name).cpu(), snapshot[name].cpu(), rtol = 1e-4, atol = 1e-6 + ), ( + f"{name} was not restored to its scaled value by loader._fix_rope_inv_freq " + "after the transformers v5 buffer blank (issue #2405 / PR #6907)." + ) + + def test_object_style_rope_scaling_does_not_crash(): # Object-style rope_scaling must be normalized, not .get()'d directly. from dataclasses import dataclass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 169b610988..1aa2c6e820 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2834,6 +2834,7 @@ def patch_llama_rope_scaling( dim = self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, + config=self.config, ) elif scaling_type == "longrope": self.rotary_emb = {longrope_rope_function}( diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c25a031b82..a1da099758 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1930,11 +1930,18 @@ class LlamaExtendedRotaryEmbedding(LlamaRotaryEmbedding): # From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41 def _apply_inv_freq_scaling(self, freqs: torch.Tensor): - # Values obtained from grid search - scale_factor = 8 - low_freq_factor = 1 - high_freq_factor = 4 - old_context_len = 8192 # original llama3 length + # llama3 factors from config; Llama-3.1 defaults when built without one + # (legacy codegen path). Hardcoding 8 is wrong for e.g. Llama-3.2 (32). + # v5 renames rope_scaling -> rope_parameters; read either so the factor + # survives even if the rope_scaling back-compat shim is dropped. + config = getattr(self, "_unsloth_rope_config", None) + rope_scaling = _rope_scaling_as_dict( + getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {} + ) + scale_factor = rope_scaling.get("factor", 8) + low_freq_factor = rope_scaling.get("low_freq_factor", 1) + high_freq_factor = rope_scaling.get("high_freq_factor", 4) + old_context_len = rope_scaling.get("original_max_position_embeddings", 8192) low_freq_wavelen = old_context_len / low_freq_factor high_freq_wavelen = old_context_len / high_freq_factor From 414503745e71b52c37117d7a96894ea46fc13ec5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:30:21 -0700 Subject: [PATCH 06/19] Run the malware gate on the RAG embedding model before it loads (#6887) * Run the malware gate on the RAG embedding model before it loads Setting the RAG embedding model through PUT /api/settings/embedding-model persisted an arbitrary repo and later handed it straight to SentenceTransformer, which deserializes pickle weights. Unlike the normal model-load paths, this route never ran evaluate_file_security, and force skipped verification entirely, so a repo Hugging Face flags as unsafe (or any repo under force) could be downloaded and loaded in the backend process without a scan. Run the malware/pickle scan at both ends: the settings endpoint now scans before persisting and returns 409 on a flagged repo even under force (force still only skips the is-embedding-model type check for offline or local repos), and the embedder scans again at the load sink so a name that arrives via env or default is covered too. Local paths and unreachable scans fail open inside evaluate_file_security, and the sink never bricks the embedder on a gate error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Thread the load token into the embedding scan and hard-fail on a block The load-sink scan ran without a token, so evaluate_file_security (which passes token=False when none is given) could not reach a gated or private repo and failed open for exactly the model SentenceTransformer would still load. Resolve the loader's own token (HF_TOKEN env or the cached login) and pass it to the sink scan, and fall back to it in the settings endpoint when the request omits one. The sink previously raised a plain RuntimeError, which the llama-server fallback in encode() and _build_st_backend_or_fallback() swallowed as a routine ST failure, silently switching backends instead of blocking. Raise a distinct UnsafeEmbeddingModelError that both fallback paths re-raise, so a flagged model hard-fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scan sentence-transformers module dirs and scope the embedding pickle gate to the ST backend Extend the RAG embedding malware gate so a poisoned pickle under a SentenceTransformer module dir (for example 0_Transformer/pytorch_model.bin) blocks. Those dirs are read from the repo's modules.json and passed as load roots to evaluate_file_security at both the settings endpoint and the load sink, so such a pickle is treated as root-level there instead of an unreferenced nested shard that was previously allowed. Scope the ST pickle scan to the sentence-transformers backend. On the llama-server backend the embedder loads GGUF files (inert) from the -GGUF companion repo, never the ST repo's pickle, so a custom ST repo with a flagged pickle and a clean GGUF companion is no longer rejected. The existing GGUF availability checks already cover that path. Return 403 for the hard security block instead of 409. The settings UI routes every 409 into the forceable save-anyway flow, but this block cannot be bypassed by force, so it now uses a distinct status the client treats as non-forceable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Base the embedding pickle scan on the actual backend, not just the resolver _llama_backend_active only consulted the auto resolver, so on a GPU box where auto resolves to sentence-transformers but the process already fell back to the llama-server backend at runtime (a torch or CUDA load/encode failure), it returned False and the settings endpoint hard-blocked a save whose ST pickle is flagged even though the process loads only inert GGUF. Add active_backend_is_llama, which reflects the actual built backend (True when the cached backend is a LlamaServerBackend, including a runtime fallback) and otherwise defers to the resolver as a fresh process would, and delegate _llama_backend_active to it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Report the cached embedding backend verbatim, not the resolver active_backend_is_llama() fell through to the config resolver whenever a backend was already built but was not llama-server, so a live sentence-transformers backend could report llama=True once the resolver picked llama (GPU heuristic or a runtime config change) and wrongly skip its pickle scan. Once a backend exists, return isinstance(backend, LlamaServerBackend) directly; only defer to the resolver before any backend is built. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/rag/embeddings.py | 117 ++++++ studio/backend/routes/settings.py | 78 +++- .../test_embedding_model_security_gate.py | 365 ++++++++++++++++++ .../tests/test_security_gate_consistency.py | 13 + .../features/settings/api/embedding-model.ts | 9 + .../features/settings/tabs/general-tab.tsx | 6 +- 6 files changed, 573 insertions(+), 15 deletions(-) create mode 100644 studio/backend/tests/test_embedding_model_security_gate.py diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 345b4dd853..47d26209b4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None: install_torchao_windows_rocm_stub() +class UnsafeEmbeddingModelError(RuntimeError): + """Raised when the embedding model repo is flagged unsafe. A distinct type so the + llama-server fallback paths re-raise it instead of masking a security block as a + routine ST failure.""" + + +def _ambient_hf_token() -> str | None: + """The HF token the loader itself would use (HF_TOKEN env or the cached login), so + the scan can reach a gated/private repo instead of failing open. None if unavailable.""" + try: + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: + """The module directories a SentenceTransformer load reads weights from, taken from + the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). + ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the + security scan: a flagged pickle directly under one must block. Returns () on any + failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + """ + try: + import json + + from utils.paths import is_local_path + + if is_local_path(name): + from pathlib import Path + from utils.paths import normalize_path + + path = Path(normalize_path(name)).expanduser() / "modules.json" + if not path.is_file(): + return () + data = json.loads(path.read_text()) + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + local = hf_hub_download(name, "modules.json", token = token or None) + except EntryNotFoundError: + return () + data = json.loads(open(local).read()) + subdirs = [] + for module in data or (): + sub = str((module or {}).get("path", "")).strip().strip("/") + if sub: + subdirs.append(sub) + return tuple(dict.fromkeys(subdirs)) + except Exception: + return () + + +def _guard_model_security(name: str) -> None: + """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside + SentenceTransformer regardless of trust_remote_code. Defense in depth behind the + /settings gate (a name can also arrive via env/default); local paths and unreachable + scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error. + """ + try: + from utils.security import evaluate_file_security, security_load_subdirs + + token = _ambient_hf_token() + # Union the audio-model load roots with the ST module dirs so a flagged pickle + # directly under a Transformer module dir (0_Transformer/) blocks instead of + # passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) + ) + blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + except Exception: + return + if blocked: + raise UnsafeEmbeddingModelError( + f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security " + "scan; refusing to load. Set a different RAG embedding model." + ) + + def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" @@ -75,6 +156,7 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) + _guard_model_security(name) _model = SentenceTransformer( name, device = device, model_kwargs = {"torch_dtype": "float16"} ) @@ -159,6 +241,8 @@ class _SentenceTransformersBackend: ): try: return _st_encode(texts, model_name = model_name, normalize = normalize) + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure # ST loaded but this encode blew up; swap the process to the llama-server # embedder (so later encodes stay in one space) and retry. @@ -222,6 +306,8 @@ def _build_st_backend_or_fallback(): try: backend.warm(model_name = None) return backend + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure fallback = _try_make_llama_backend() if fallback is None: @@ -290,6 +376,37 @@ def _reset_backend() -> None: _backend_key = None +def active_backend_is_llama() -> bool: + """True when this process actually embeds via the llama-server (GGUF) backend. + + Reflects the ACTUAL built backend once one exists: an ``auto`` install that + resolves to sentence-transformers but then falls back to llama-server at + runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or + ``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so + callers gating on the ST pickle must see llama here. Before any backend is + built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw + key) exactly as a fresh process would. Never raises: a backend probe must not + block saving a model.""" + try: + with _backend_lock: + backend = _backend + if backend is not None: + # A backend exists: report what it ACTUALLY is. A concrete + # sentence-transformers backend must return False even if the + # resolver would now pick llama, so its pickle stays gated. If the + # llama import fails we cannot be llama, so fall to the safe False. + try: + from .embed_llama_server import LlamaServerBackend + except Exception: # noqa: BLE001 - llama plumbing import must never block + return False + return isinstance(backend, LlamaServerBackend) + raw = (config.EMBED_BACKEND or "auto").strip().lower() + key = _resolve_auto() if raw in _AUTO_ALIASES else raw + return key in _LLAMA_ALIASES + except Exception: # noqa: BLE001 - a backend probe must never block saving + return False + + def warm(model_name: str | None = None) -> None: """Eagerly load the embedder so the first real request isn't slow.""" _get_backend().warm(model_name = model_name) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 862bce8be8..914699f540 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -260,17 +260,29 @@ def _embedding_model_response() -> EmbeddingModelResponse: ) -def _llama_backend_active() -> bool: - """True when this install embeds via the llama-server (GGUF) backend.""" - from core.rag import config as rag_config - from core.rag import embeddings - +def _ambient_hf_token() -> Optional[str]: + """The HF token the loader would use (HF_TOKEN env or the cached login), so a gated + repo is scanned rather than failing open. None if unavailable.""" try: - raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() - key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _llama_backend_active() -> bool: + """True when this install actually embeds via the llama-server (GGUF) backend. + + Delegates to the embeddings module so a runtime fallback from + sentence-transformers to llama-server (after a torch/CUDA load or encode + failure) is honored: in that state the process loads only inert GGUF, so the + ST pickle gate below must not hard-block a repo whose GGUF companion is clean. + Before any backend is built this still reflects the resolver.""" + from core.rag import embeddings + try: + return embeddings.active_backend_is_llama() except Exception: # noqa: BLE001 - backend probe must never block saving return False - return key in embeddings._LLAMA_ALIASES def _resolves_as_local_gguf(model: str) -> bool: @@ -357,6 +369,8 @@ def update_embedding_model( """Set the RAG embedding model. Unless ``force`` is set, the repo is verified to be an embedding model via HF metadata; an unverifiable model (wrong type, typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + A repo flagged unsafe by HF's security scan returns 403 instead: a hard block + that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" from utils.models import is_embedding_model @@ -370,15 +384,51 @@ def update_embedding_model( event = "settings.update_embedding_model_failed", log = logger, ) from exc + hf_token = (payload.hf_token or "").strip() or None # The env/default model needs no verification; saving it is a no-op override. # A local GGUF on the llama-server backend is accepted as-is: it is exactly # what the backend loads, and HF metadata cannot verify a local path. - if ( - model != default_embedding_model() - and not payload.force - and not (_llama_backend_active() and _resolves_as_local_gguf(model)) - ): - hf_token = (payload.hf_token or "").strip() or None + is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model) + # The pickle gate only matters for the sentence-transformers backend, which is what + # deserializes pickles. On the llama-server backend the embedder loads GGUF files + # (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would + # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability + # checks below cover that path instead. + scan_st_pickle = ( + model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() + ) + if scan_st_pickle: + # Malware/pickle gate before we persist a repo the embedder later loads with + # SentenceTransformer. Runs even under force (force only skips the is-embedding + # type check for offline/local repos HF cannot verify); local paths and + # unreachable scans fail open inside evaluate_file_security. + from utils.security import evaluate_file_security, security_load_subdirs + from core.rag.embeddings import _st_module_subdirs + + # Fall back to the loader's own token so a gated/private repo is actually scanned + # (a token-less scan fails open for exactly the repo that would still load). + scan_token = hf_token or _ambient_hf_token() + # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under + # one blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + ( + *security_load_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token), + ) + ) + ) + if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + # 403, not 409: the client routes every 409 into the forceable "save anyway" + # flow, but this block is a hard, non-forceable security refusal. + raise HTTPException( + status_code = 403, + detail = ( + f"{model!r} is flagged as unsafe by Hugging Face's security scan and " + "cannot be used as the embedding model." + ), + ) + if model != default_embedding_model() and not payload.force and not is_local_gguf: from core.rag import config as rag_config # A GGUF-named repo on the llama-server backend is loaded from its .gguf diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py new file mode 100644 index 0000000000..940b35d7ba --- /dev/null +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The RAG embedding model must pass the malware/pickle gate before it is persisted or +loaded. A flagged repo (or any repo saved with force) previously reached +SentenceTransformer unscanned, bypassing the normal model-load protections.""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.settings as settings + + +class _Decision: + def __init__(self, blocked): + self.blocked = blocked + + +def _security_stub(blocked): + mod = _types.ModuleType("utils.security") + mod.evaluate_file_security = lambda *a, **k: _Decision(blocked) + mod.security_load_subdirs = lambda *a, **k: () + return mod + + +@pytest.fixture +def client(monkeypatch): + # The settings scan unions in the ST module dirs read from modules.json; keep it + # offline and deterministic for the endpoint tests that use this fixture. + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + return TestClient(app, raise_server_exceptions = False), saved + + +def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put( + "/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True} + ) + # 403, not the forceable 409, so the client does not offer "save anyway". + assert r.status_code == 403 + assert "model" not in saved # force must not persist a flagged repo + + +def test_flagged_repo_is_blocked_without_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert r.status_code == 403 + assert "model" not in saved + + +def test_hard_block_uses_non_forceable_status(client, monkeypatch): + # The forceable verification path uses 409; the hard security block must be distinct + # (403) so the frontend never routes it into the "save anyway" force flow. + c, _saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert blocked.status_code == 403 + + # A verification failure (not-an-embedding-model) stays forceable at 409. + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False) + import utils.models as _models + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"}) + assert unverified.status_code == 409 + + +def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): + # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's + # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + # force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped. + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama path + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): + # auto resolves to sentence-transformers (GPU present) but the embedder fell back to + # llama-server at runtime (torch/CUDA load or encode failure), so the process now loads + # only inert GGUF. The real _llama_backend_active() must reflect that cached fallback, + # so a flagged ST repo with a clean GGUF companion must not be hard-blocked here. + import core.rag.embeddings as embeddings + from core.rag.embed_llama_server import LlamaServerBackend + + # Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even + # though the auto resolver would still say sentence-transformers. + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + # Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the + # real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored. + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch): + # active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers + # to the resolver (fresh-process behavior) when none has been built yet. + import core.rag.embeddings as embeddings + import core.rag.config as rag_config + from core.rag.embed_llama_server import LlamaServerBackend + + # A cached llama backend wins even when auto would resolve to sentence-transformers. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto") + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + assert embeddings.active_backend_is_llama() is True + + # A cached ST backend reports False even when the resolver now picks llama, so its + # pickle stays gated (the cached backend, not the resolver, is what actually embeds). + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend()) + assert embeddings.active_backend_is_llama() is False + + # No cached backend -> the resolver decides, unchanged from before. + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", None) + assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers + + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + assert embeddings.active_backend_is_llama() is True # auto -> llama-server + + # An explicit (non-auto) key is honored verbatim without a cached backend. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server") + assert embeddings.active_backend_is_llama() is True + + +def test_settings_scan_scopes_module_subdirs(monkeypatch): + # The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a + # pickle directly under one blocks; assert those subdirs reach evaluate_file_security. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + import core.rag.embeddings as embeddings + + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda *a, **k: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True} + ) + assert r.status_code == 200 + assert "0_Transformer" in seen["subdirs"] + + +def test_clean_repo_saves_under_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "acme/clean-embed" + + +def test_load_sink_refuses_flagged_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + import core.rag.embeddings as embeddings + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._guard_model_security("attacker/malicious-embed") + + +def test_load_sink_allows_clean_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import core.rag.embeddings as embeddings + embeddings._guard_model_security("acme/clean-embed") # no raise + + +def test_sink_threads_ambient_token_into_scan(monkeypatch): + # A gated repo set via env/default has no request token; the guard must feed the + # loader's own token to the scan, or it fails open for the repo that still loads. + seen = {} + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = ( + lambda name, token = None: seen.setdefault("subdirs_token", token) or () + ) + mod.evaluate_file_security = lambda *a, **k: seen.setdefault( + "scan_token", k.get("hf_token") + ) or _Decision(False) + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient") + embeddings._guard_model_security("acme/gated-embed") + assert seen["scan_token"] == "hf_ambient" + assert seen["subdirs_token"] == "hf_ambient" + + +def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch): + # A flagged pickle directly under a Transformer module dir (0_Transformer/) must + # reach the scan as a load root; assert the guard unions the module dirs into + # load_subdirs so evaluate_file_security treats such a pickle as root-level. + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda name, token = None: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + embeddings._guard_model_security("acme/embed-with-module-dir") + assert "0_Transformer" in seen["subdirs"] + + +def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch): + # The helper must parse each module's non-empty "path" from a local repo's + # modules.json and drop the root-level ("") Transformer entry. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [ + {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, + {"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."}, + {"idx": 2, "name": "2", "path": "", "type": "..."}, + ] + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None) + assert subdirs == ("0_Transformer", "1_Pooling") + + +def test_st_module_subdirs_swallows_errors(monkeypatch): + # Any failure (no modules.json, offline, malformed) returns () so the guard never + # bricks the embedder. + import huggingface_hub + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise RuntimeError("offline") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom) + assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == () + + +def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): + # The ST encode fallback must re-raise a security block, not swap to llama-server. + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise embeddings.UnsafeEmbeddingModelError("flagged") + + monkeypatch.setattr(embeddings, "_st_encode", _boom) + monkeypatch.setattr( + embeddings, + "_switch_to_llama_fallback", + lambda err: pytest.fail("security block must not fall back to llama-server"), + ) + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._SentenceTransformersBackend().encode(["hi"]) diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index db66df8a30..b5f1069f12 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base(): if runs_gate and not resolves_base: offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") assert not offenders, "\n".join(offenders) + + +def test_rag_embedding_path_runs_the_malware_gate(): + """The RAG embedding model is set through /settings and later loaded by + SentenceTransformer, which deserializes pickles; both sites must run the malware gate + or a flagged repo loads unscanned (bypassing the normal model-load protections).""" + offenders = [] + for rel in ("routes/settings.py", "core/rag/embeddings.py"): + if "evaluate_file_security(" not in (_BACKEND / rel).read_text(): + offenders.append( + f"{rel} loads/persists an embedding model without evaluate_file_security" + ) + assert not offenders, "\n".join(offenders) diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts index 8b6bc7ee7f..9a61142f73 100644 --- a/studio/frontend/src/features/settings/api/embedding-model.ts +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = { * (wrong type, gated repo, or offline). Retry with force to save anyway. */ export class EmbeddingModelVerificationError extends Error {} +/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan. + * A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */ +export class EmbeddingModelBlockedError extends Error {} + function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { return { embeddingModel: settings.embedding_model, @@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings( force: options?.force ?? false, }), }); + if (res.status === 403) { + throw new EmbeddingModelBlockedError( + await readFastApiError(res, "This model is blocked by a security scan"), + ); + } if (res.status === 409) { throw new EmbeddingModelVerificationError( await readFastApiError(res, "Could not verify the embedding model"), diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 7670aae5fa..8fd70cc1b7 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -42,6 +42,7 @@ import { updatePreviewSharing, } from "../api/preview-sharing"; import { + EmbeddingModelBlockedError, type EmbeddingModelSettings, EmbeddingModelVerificationError, loadEmbeddingModelSettings, @@ -410,7 +411,10 @@ export function GeneralTab() { description: t("settings.general.rag.reindexWarning"), }); } catch (error) { - if (error instanceof EmbeddingModelVerificationError) { + // A hard security block cannot be forced; keep the "save anyway" action hidden. + if (error instanceof EmbeddingModelBlockedError) { + setEmbeddingModelNeedsForce(false); + } else if (error instanceof EmbeddingModelVerificationError) { setEmbeddingModelNeedsForce(true); } setEmbeddingModelError( From d79495dc96cf7c6ab0d60585e232fa381549c14e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:41:37 -0700 Subject: [PATCH 07/19] Add RDNA 2/3/4 ROCm routing tests via a CPU-only torch spoof (#6935) * Add RDNA 2/3/4 ROCm routing tests via a CPU-only torch spoof Introduces tests/_zoo_rocm_spoof.py, the ROCm sibling of _zoo_aggressive_cuda_spoof.py: it reuses the CUDA spoof's torch.cuda no-op machinery and overlays an AMD Radeon identity (torch.version.hip, gcnArchName, capability) for any RDNA 2/3/4 gfx target, so hip code paths run on CPU-only CI with no AMD hardware. tests/studio/install/test_rocm_rdna_routing.py then asserts unsloth_zoo routes every RDNA arch (gfx1030/1031/1032/1034, gfx1100/1101/1102, gfx1150/1151, gfx1200/1201) correctly: device_type resolves to hip, llama.cpp target resolves to (rocm, gfx), and the per-family ROCm bundle suffix (gfx103X/gfx110X/gfx120X, or self for gfx1150/1151) is picked. The torch-facing checks run in a subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached at import) resolves from a clean process; the pure gfx-family mapping runs in-process. Guarded by importorskip so it runs where torch and unsloth_zoo are installed (the Repo tests CPU job) and skips elsewhere. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/_zoo_rocm_spoof.py | 84 +++++++++++++++++++ .../studio/install/test_rocm_rdna_routing.py | 84 +++++++++++++++++++ 2 files changed, 168 insertions(+) create mode 100644 tests/_zoo_rocm_spoof.py create mode 100644 tests/studio/install/test_rocm_rdna_routing.py diff --git a/tests/_zoo_rocm_spoof.py b/tests/_zoo_rocm_spoof.py new file mode 100644 index 0000000000..050191e9d1 --- /dev/null +++ b/tests/_zoo_rocm_spoof.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""ROCm/RDNA spoof: present torch as an AMD Radeon (RDNA 2/3/4) card on a +GPU-less host, so hip paths (device_type -> "hip", llama.cpp ROCm bundle) are +testable in CPU-only CI with no AMD hardware. The ROCm sibling of +_zoo_aggressive_cuda_spoof.py: it reuses that spoof's torch.cuda no-op machinery +and overlays the AMD identity (torch.version.hip, gcnArchName, Radeon name). +Apply BEFORE importing unsloth/unsloth_zoo, since DEVICE_TYPE is cached there. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys + +# gfx -> (marketing name, (capability major, minor), torch.version.hip). hip is +# the ROCm build torch was made against (RDNA2/3 ship 6.x; gfx1102/115x/RDNA4 7.2). +_PROFILES: dict[str, tuple[str, tuple[int, int], str]] = { + "gfx1030": ("AMD Radeon RX 6900 XT", (10, 3), "6.4.43483"), # RDNA2 + "gfx1031": ("AMD Radeon RX 6700 XT", (10, 3), "6.4.43483"), + "gfx1032": ("AMD Radeon RX 6600", (10, 3), "6.4.43483"), + "gfx1034": ("AMD Radeon RX 6400", (10, 3), "6.4.43483"), + "gfx1100": ("AMD Radeon RX 7900 XTX", (11, 0), "6.4.43483"), # RDNA3 + "gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"), + "gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"), + "gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU + "gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"), + "gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4 + "gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"), +} + + +def _cuda_spoof(): + """Load the sibling CUDA spoof by path (robust to sys.path), so we reuse its + torch.cuda machinery instead of duplicating it.""" + if "_zoo_aggressive_cuda_spoof" in sys.modules: + return sys.modules["_zoo_aggressive_cuda_spoof"] + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py") + spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + sys.modules["_zoo_aggressive_cuda_spoof"] = mod + return mod + + +def apply(gfx: str = "gfx1100", device_count: int = 1) -> None: + """Present torch as `gfx`. Re-callable to switch arch (identity is overlaid; + the underlying no-op machinery is applied once).""" + import torch + + if gfx not in _PROFILES: + raise KeyError(f"Unknown gfx {gfx!r}; known: {', '.join(_PROFILES)}") + name, cap, hip = _PROFILES[gfx] + + _cuda_spoof().apply() # is_available/device_count/streams/rng/amp/... + + # Overlay the AMD identity on top of the (NVIDIA-shaped) CUDA spoof. + torch.version.hip = hip + torch.version.cuda = None + torch.cuda.device_count = lambda: device_count + torch.cuda.get_device_name = lambda *a, **k: name + torch.cuda.get_device_capability = lambda *a, **k: cap + torch.cuda.get_arch_list = lambda: [gfx] + + class _Props: + pass + + _p = _Props() + _p.name = name + _p.gcnArchName = f"{gfx}:sramecc-:xnack-" # ROCm advertises feature flags + _p.major, _p.minor = cap + _p.total_memory = 16 * 1024**3 + _p.multi_processor_count = 40 + _p.warp_size = 32 # RDNA wavefront (CDNA is 64) + _p.is_integrated = gfx in ("gfx1150", "gfx1151") + _p.is_multi_gpu_board = False + torch.cuda.get_device_properties = lambda *a, **k: _p + + +if __name__ == "__main__": + apply() + import torch + print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName) diff --git a/tests/studio/install/test_rocm_rdna_routing.py b/tests/studio/install/test_rocm_rdna_routing.py new file mode 100644 index 0000000000..b4aeafb7e4 --- /dev/null +++ b/tests/studio/install/test_rocm_rdna_routing.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""RDNA 2/3/4 routing, validated on CPU-only CI with no AMD hardware. + +tests/_zoo_rocm_spoof.py presents torch as each Radeon gfx arch, then we assert +unsloth_zoo routes it: device_type -> "hip", llama.cpp target -> ("rocm", gfx), +and the per-family ROCm bundle suffix. The torch-facing checks run in a +subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached +at import) resolves from a clean process. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("unsloth_zoo") + +_TESTS_DIR = Path(__file__).resolve().parents[2] # tests/ + +# gfx -> (expected llama.cpp target, expected ROCm bundle family). +_ARCHES = { + "gfx1030": (("rocm", "gfx1030"), "gfx103X"), # RDNA2 + "gfx1031": (("rocm", "gfx1031"), "gfx103X"), + "gfx1032": (("rocm", "gfx1032"), "gfx103X"), + "gfx1034": (("rocm", "gfx1034"), "gfx103X"), + "gfx1100": (("rocm", "gfx1100"), "gfx110X"), # RDNA3 + "gfx1101": (("rocm", "gfx1101"), "gfx110X"), + "gfx1102": (("rocm", "gfx1102"), "gfx110X"), + "gfx1150": (("rocm", "gfx1150"), "gfx1150"), # RDNA3.5 APU (self-family) + "gfx1151": (("rocm", "gfx1151"), "gfx1151"), + "gfx1200": (("rocm", "gfx1200"), "gfx120X"), # RDNA4 + "gfx1201": (("rocm", "gfx1201"), "gfx120X"), +} + +# Child: spoof each arch, then record device_type once (fresh import) and the +# live llama.cpp target per arch. Emits one JSON line the parent parses. +_CHILD = """ +import json, sys +sys.path.insert(0, {tests!r}) +import _zoo_rocm_spoof as spoof +arches = {arches!r} +spoof.apply(arches[0]) +from unsloth_zoo.device_type import get_device_type, is_hip +device_type = [get_device_type(), is_hip()] +from unsloth_zoo import llama_cpp as lc +targets = {{}} +for gfx in arches: + spoof.apply(gfx) + targets[gfx] = list(lc._detect_gpu_target()) +print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}})) +""" + + +@pytest.fixture(scope = "module") +def routed(): + code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES)) + proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True) + line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None) + assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + return json.loads(line[len("RESULT ") :]) + + +@pytest.mark.parametrize("gfx", list(_ARCHES)) +def test_detect_gpu_target(routed, gfx): + # RDNA card is routed to its ROCm gfx target (drives the llama.cpp bundle). + assert tuple(routed["targets"][gfx]) == _ARCHES[gfx][0] + + +def test_device_type_is_hip(routed): + # An RDNA card must resolve the compute device_type to "hip". + assert routed["device_type"] == ["hip", True] + + +@pytest.mark.parametrize("gfx", list(_ARCHES)) +def test_rocm_gfx_family(gfx): + # Pure mapping (no torch): each gfx picks the right per-family ROCm bundle. + from unsloth_zoo import llama_cpp as lc + assert lc._rocm_gfx_family(gfx) == _ARCHES[gfx][1] From 59977f95c318c1ba81b53c5b50d4a0ef9c342fea Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 05:49:24 -0700 Subject: [PATCH 08/19] GRPO: default router_aux_loss_coef to 0 on TRL >= 1.7.0 (#6938) TRL 1.7.0 enables the MoE router load-balancing aux loss by default (router_aux_loss_coef = 0.001). Unsloth's optimized GRPO forward does not compute it, so default the coefficient to 0, matching pre-1.7.0 behaviour. Users can still opt in with router_aux_loss_coef > 0. No-op on TRL < 1.7.0. --- unsloth/models/rl.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 62ef9e916a..b5cadf2dea 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1370,6 +1370,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): # [TODO] See https://fengyao.notion.site/off-policy-rl # https://github.com/huggingface/trl/pull/3867 (August 7th) "vllm_importance_sampling_correction": False, + # TRL >= 1.7.0 enables the MoE router aux loss by default (0.001); the optimized + # GRPO forward does not compute it, so default off. Opt in via router_aux_loss_coef > 0. + "router_aux_loss_coef": 0.0, } for k, v in replacements.items(): x = f"{k}( = [^,\n]{{1,}})?,\n" From 411c4d1e50362c0fe6c27d4eea0c29171f79a27a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:13:43 -0700 Subject: [PATCH 09/19] Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning (#6908) * Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning Adds unsloth/DeepSeek-V4-Flash-GGUF as a default selectable model with the recommended decoding defaults (temperature 1.0, top_p 1.0 from the official generation_config.json) and its three tier reasoning control. The high/max ladder is surfaced for deepseek-v4 model ids and flows through the existing enable_thinking_effort reasoning style via chat_template_kwargs, so no frontend changes are needed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio DeepSeek-V4: segment-scope high, enable thinking for lone effort, render tests Match deepseek-v4 on whole repo-name segments so a future deepseek-v40 or deepseek40 cannot false-match the synthetic 'high'. In _request_reasoning_kwargs, emit enable_thinking when a named effort level is sent without it, so the newly exposed High mode renders thinking-on over the API (the UI already sent it explicitly). Add a none/high/max render-path test file (jinja behind importorskip) with a lone-high regression. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../assets/configs/inference_defaults.json | 9 +- studio/backend/core/inference/defaults.py | 2 + studio/backend/core/inference/llama_cpp.py | 18 +- .../tests/test_deepseek_v4_thinking_effort.py | 181 ++++++++++++++++++ .../test_safetensors_capability_advertise.py | 38 ++++ 5 files changed, 245 insertions(+), 3 deletions(-) create mode 100644 studio/backend/tests/test_deepseek_v4_thinking_effort.py diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1c7a409bc1..0633f80bbc 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -235,6 +235,13 @@ "min_p": 0.1, "repetition_penalty": 1.0 }, + "deepseek-v4": { + "temperature": 1.0, + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "repetition_penalty": 1.0 + }, "deepseek-r1": { "temperature": 0.6, "top_p": 0.95, @@ -394,7 +401,7 @@ "phi-4", "phi-3", "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral", "devstral", "pixtral", - "deepseek-r1", "deepseek-v3", "deepseek-ocr", + "deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", "minimax-m2.7", "minimax-m2.5", "minimax", diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index b64605e16f..a1d03c03e0 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -8,6 +8,7 @@ import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", @@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [ DEFAULT_MODELS_STANDARD = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5e6287f528..8b40f5fccd 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -686,6 +686,16 @@ def detect_reasoning_flags( else [] ) if effort_levels: + # DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its + # template only branches on 'max', so the literal scan misses 'high'. Add it + # (matched on whole repo-name segments, so 'deepseek-v40' won't false-match) + # to expose the full none/high/max ladder instead of none/max. + segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1]) + is_dsv4 = "deepseek4" in segments or any( + a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:]) + ) + if is_dsv4 and "high" not in effort_levels: + effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index) # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -1741,9 +1751,13 @@ class LlamaCppBackend: # 'low' effort the way gpt-oss does (those models genuinely # cannot disable). thinking_off = enable_thinking is False or reasoning_effort == "none" - if enable_thinking is not None or reasoning_effort == "none": + # A named effort level implies thinking on, so emit enable_thinking + # even if the caller sent only reasoning_effort (else the template + # defaults it off and the requested level never renders). + effort_on = reasoning_effort in self._reasoning_effort_levels + if enable_thinking is not None or reasoning_effort == "none" or effort_on: kwargs["enable_thinking"] = not thinking_off - if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + if not thinking_off and effort_on: kwargs["reasoning_effort"] = reasoning_effort elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py new file mode 100644 index 0000000000..19808ad0d7 --- /dev/null +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""DeepSeek-V4-Flash reasoning toggle: None / High / Max. + +The GGUF template gates thinking with ``enable_thinking`` and only branches +``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking). +Detection used to return the single level ``['max']``, so the UI collapsed to +None / Max and the plain-thinking tier was unreachable. Detection now surfaces +``'high'`` as that plain tier, giving None / High / Max. These tests pin the +classifier, the GLM-style parity case, and the full request-kwargs -> rendered +prompt path for each state (the model itself is too large to load here). +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking +# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think +# fallback. Any non-'max' effort renders as ordinary thinking. +DEEPSEEK_V4_TEMPLATE = """ +{%- if not thinking is defined -%} + {%- if enable_thinking is defined -%} + {%- set thinking = enable_thinking -%} + {%- else -%} + {%- set thinking = false -%} + {%- endif -%} +{%- endif -%} +{%- if not reasoning_effort is defined -%} + {%- set reasoning_effort = none -%} +{%- endif -%} +{{- bos_token -}} +{%- if thinking and reasoning_effort == 'max' -%} + {{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}} +{%- endif -%} +{%- for message in messages -%} + {{- '<|User|>' + (message['content'] or '') -}} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|Assistant|>' -}} + {%- if thinking -%}{{- '' -}}{%- else -%}{{- '' -}}{%- endif -%} +{%- endif -%} +""" + + +# GLM-5.2-style: branches on two effort literals, so 'high' already exists as +# the sub-'max' tier and detection must leave the pair untouched. +GLM_STYLE_TEMPLATE = """ +{%- if enable_thinking -%} + {%- if reasoning_effort == 'high' -%}{{- 'H' -}} + {%- elif reasoning_effort == 'max' -%}{{- 'M' -}} + {%- endif -%} +{%- endif -%} +""" + + +# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped +# to deepseek-v4, so this must stay ['max'] (no phantom 'high'). +NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE + + +# A template whose sole effort literal is a sub-'max' level: the guard targets +# only the ['max']-alone case, so a lone 'high' stays a singleton. +HIGH_ONLY_TEMPLATE = """ +{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%} +""" + + +def _render(template: str, **kwargs) -> str: + jinja2 = pytest.importorskip("jinja2") + env = jinja2.Environment() + tmpl = env.from_string(template) + return tmpl.render(bos_token = "", add_generation_prompt = True, **kwargs) + + +# -- Classifier ------------------------------------------------------- + + +def test_deepseek_v4_surfaces_high_as_plain_tier(): + """Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_glm_style_two_level_template_unchanged(): + """A template that already names a sub-'max' tier is left as-is.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_synthetic_high_scoped_to_deepseek_v4(): + """The same ['max']-only template under a non-deepseek id keeps ['max'].""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_effort_levels"] == ["max"] + + +def test_guard_does_not_fire_for_sub_max_singleton(): + """The expansion targets only ['max']; a lone 'high' stays a singleton.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only") + assert flags["reasoning_effort_levels"] == ["high"] + + +# -- Request kwargs -> rendered prompt, for each state ---------------- + + +def _kwargs_for(flags: dict, enable_thinking, reasoning_effort): + """Drive the real backend method with a shim carrying the detected flags.""" + from core.inference.llama_cpp import LlamaCppBackend + + shim = SimpleNamespace( + _supports_reasoning = flags["supports_reasoning"], + _reasoning_always_on = flags["reasoning_always_on"], + _reasoning_style = flags["reasoning_style"], + _reasoning_effort_levels = flags["reasoning_effort_levels"], + _supports_preserve_thinking = flags["supports_preserve_thinking"], + ) + build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim) + return build(enable_thinking, reasoning_effort, None) or {} + + +def _flags(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + + +def test_none_state_renders_non_thinking(): + """UI 'None' -> enable_thinking=false -> closed , no preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None) + assert kwargs == {"enable_thinking": False} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_high_state_renders_plain_thinking(): + """UI 'High' -> et=true, effort=high -> open , no max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_max_state_injects_max_preamble(): + """UI 'Max' -> et=true, effort=max -> open plus the max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" in out + + +def test_high_effort_alone_enables_thinking(): + """API caller sending only reasoning_effort='high' (no enable_thinking) still + gets thinking on, so the newly exposed High mode renders correctly.""" + kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 9fd1535f22..0ed670ac01 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -48,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }} """ +# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort +# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders +# identically to thinking-on-without-the-preamble), so the literal scan alone +# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to +# expose the encoder's full none/high/max ladder. +DEEPSEEK_V4_TEMPLATE = ( + "{%- if not thinking is defined %}" + "{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}" + "{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n" + "{%- if thinking and reasoning_effort == 'max' %}" + "{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n" + "{%- for message in messages %}{{- message.content }}{%- endfor %}" +) + + PLAIN_TEMPLATE = """ {%- for message in messages %} {{- message.role + ': ' + message.content + '\\n' }} @@ -90,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false(): assert flags["reasoning_style"] == "enable_thinking" +def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max(): + """DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble. + Classified as the hybrid style with the full none/high/max ladder even + though the template only branches on 'max'.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected(): + """The 'high' injection is scoped to deepseek-v4: a different model whose + template only branches on 'max' keeps ['max'] (no phantom 'high').""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["max"] + + def test_detect_safetensors_features_passes_template_through_to_classifier(): """Route wrapper forwards a real template to the inner classifier.""" from routes.inference import _detect_safetensors_features From 10d8f985a270cd4cacf2518e9e895874da201f3e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:19:25 -0700 Subject: [PATCH 10/19] Versioning --- pyproject.toml | 6 +++--- unsloth/models/_utils.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 844ead2454..80b3d757e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1aa2c6e820..1c75f8ce66 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.6.9" +__version__ = "2026.7.1" __all__ = [ "SUPPORTS_BFLOAT16", From ba450b437eb34ee476df7ffb61bb25db365c7353 Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:26:00 +0300 Subject: [PATCH 11/19] Studio: add assistant response details panel (#6842) * Studio: add assistant response details panel * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hide model badge by default, show on hover/focus Wrap MessageResponseModelBadge in a span with hidden/group-hover visibility classes to reduce visual clutter. The badge now only displays when hovering or focusing on the assistant message, improving the UI presentation. Updated corresponding tests to verify the new CSS classes. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../message-response-details-sheet.tsx | 483 ++++++++++++++++++ .../src/components/assistant-ui/reasoning.tsx | 10 +- .../src/components/assistant-ui/thread.tsx | 118 +++-- .../src/features/chat/api/chat-adapter.ts | 92 +++- studio/frontend/src/features/chat/index.ts | 7 +- .../chat/stores/chat-preferences-store.ts | 7 + .../src/features/settings/tabs/chat-tab.tsx | 15 + .../test_chat_response_details_ui_contract.py | 93 ++++ 8 files changed, 776 insertions(+), 49 deletions(-) create mode 100644 studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx create mode 100644 tests/studio/test_chat_response_details_ui_contract.py diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx new file mode 100644 index 0000000000..823696693a --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -0,0 +1,483 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { + Sheet, + SheetContent, + SheetDescription, + SheetHeader, + SheetTitle, +} from "@/components/ui/sheet"; +import { + customProviderDisplayName, + parseExternalModelId, + useChatPreferencesStore, + useChatRuntimeStore, + useExternalProvidersStore, +} from "@/features/chat"; +import { cn } from "@/lib/utils"; +import { FileDatabaseIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useMessage, useMessageTiming } from "@assistant-ui/react"; +import type { FC, ReactNode } from "react"; + +type ResponseDetailsMetadata = { + modelId?: string; + modelLabel?: string; + responseModelId?: string; + providerId?: string; + providerName?: string; + providerType?: string; + startedAt?: number; + finishedAt?: number; + durationMs?: number; + sessionId?: string | null; + cancelId?: string; + toolCalls?: string[]; + tools?: Record; +}; + +type ContextUsageMetadata = { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + cachedTokens?: number; + cacheWriteTokens?: number; + modelId?: string; +}; + +type MessageCustomMetadata = { + responseDetails?: ResponseDetailsMetadata; + contextUsage?: ContextUsageMetadata; + serverTimings?: Record; + reasoningDuration?: number; +}; + +function asNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) + ? value + : undefined; +} + +function formatNumber(value: number | undefined): string | null { + return value == null ? null : value.toLocaleString(); +} + +function formatMs(value: number | undefined): string | null { + if (value == null) return null; + if (value < 1000) return `${Math.round(value)}ms`; + return `${(value / 1000).toFixed(2)}s`; +} + +function formatRate(value: number | undefined): string | null { + if (value == null) return null; + return `${value.toFixed(1)} tok/s`; +} + +function formatDate(value: Date | number | string | undefined): string | null { + if (value == null) return null; + const date = value instanceof Date ? value : new Date(value); + if (Number.isNaN(date.getTime())) return null; + return new Intl.DateTimeFormat(undefined, { + dateStyle: "medium", + timeStyle: "medium", + }).format(date); +} + +const TOOL_CATEGORY_LABELS: Record = { + search: "Search", + fetch: "Fetch", + code: "Code", + images: "Images", + mcp: "MCP", + docs: "Docs", + artifacts: "Canvas", +}; + +const TOOL_CALL_LABELS: Record = { + web_search: "Search", + web_fetch: "Fetch", + code_execution: "Code", + python: "Python", + terminal: "Terminal", + image_generation: "Images", + search_knowledge_base: "Docs", + render_html: "Canvas", +}; + +function uniqueValues(values: string[]): string[] { + return Array.from(new Set(values)); +} + +function toolCategoryFromCall(toolName: string): string | null { + const normalized = toolName.toLowerCase(); + if (normalized === "web_search") return "search"; + if (normalized === "web_fetch") return "fetch"; + if ( + normalized === "code_execution" || + normalized === "python" || + normalized === "terminal" + ) { + return "code"; + } + if (normalized === "image_generation") return "images"; + if (normalized === "search_knowledge_base") return "docs"; + if (normalized === "render_html") return "artifacts"; + if (normalized.startsWith("mcp__")) return "mcp"; + return null; +} + +function formatToolCallName(toolName: string): string { + const normalized = toolName.toLowerCase(); + if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized]; + if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`; + return toolName + .replace(/[_-]+/g, " ") + .replace(/\b\w/g, (letter) => letter.toUpperCase()); +} + +function toolCallsFromContent(content: unknown): string[] { + if (!Array.isArray(content)) return []; + return uniqueValues( + content + .map((part) => + part && typeof part === "object" && "type" in part + ? (part as { type?: unknown; toolName?: unknown }) + : null, + ) + .filter( + (part): part is { type: "tool-call"; toolName: string } => + part?.type === "tool-call" && + typeof part.toolName === "string" && + part.toolName.length > 0, + ) + .map((part) => part.toolName), + ); +} + +function enabledTools( + tools: Record | undefined, + toolCalls: string[], +): string | null { + if (!tools && toolCalls.length === 0) return null; + const activeKeys = new Set(); + for (const key of Object.keys(TOOL_CATEGORY_LABELS)) { + if (tools?.[key] === true) activeKeys.add(key); + } + for (const toolName of toolCalls) { + const key = toolCategoryFromCall(toolName); + if (key) activeKeys.add(key); + } + const active = Object.keys(TOOL_CATEGORY_LABELS) + .filter((key) => activeKeys.has(key)) + .map((key) => TOOL_CATEGORY_LABELS[key]); + return active.length > 0 ? active.join(", ") : "None"; +} + +function calledTools(toolCalls: string[]): string | null { + if (toolCalls.length === 0) return null; + return uniqueValues(toolCalls.map(formatToolCallName)).join(", "); +} + +function DetailSection({ + title, + children, +}: { + title: string; + children: ReactNode; +}) { + return ( +
+

{title}

+
{children}
+
+ ); +} + +function DetailRow({ + label, + value, + mono = false, +}: { + label: string; + value: ReactNode | null | undefined; + mono?: boolean; +}) { + if (value == null || value === "") return null; + return ( +
+ {label} + + {value} + +
+ ); +} + +function useResponseModelDisplay() { + const message = useMessage(); + const models = useChatRuntimeStore((s) => s.models); + const providers = useExternalProvidersStore((s) => s.providers); + + const custom = ( + message.metadata as Record | undefined + )?.custom as MessageCustomMetadata | undefined; + const responseDetails = custom?.responseDetails; + const usage = custom?.contextUsage; + const serverTimings = custom?.serverTimings; + + const recordedModelId = + responseDetails?.responseModelId ?? + responseDetails?.modelId ?? + usage?.modelId; + const parsedExternal = parseExternalModelId(recordedModelId); + const provider = parsedExternal + ? providers.find((candidate) => candidate.id === parsedExternal.providerId) + : null; + const modelSummary = models.find( + (candidate) => candidate.id === recordedModelId, + ); + const modelLabel = + responseDetails?.modelLabel ?? + responseDetails?.responseModelId ?? + parsedExternal?.modelId ?? + modelSummary?.name ?? + recordedModelId ?? + "Not recorded"; + const providerLabel = + responseDetails?.providerName ?? + provider?.name ?? + (responseDetails?.providerType + ? customProviderDisplayName(responseDetails.providerType) + : parsedExternal + ? customProviderDisplayName(provider?.providerType) + : recordedModelId + ? "Local model" + : null); + + return { + message, + custom, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + }; +} + +export const MessageResponseModelBadge: FC<{ className?: string }> = ({ + className, +}) => { + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const { modelLabel, providerLabel } = useResponseModelDisplay(); + + if (!showResponseModel || modelLabel === "Not recorded") { + return null; + } + + return ( + + {modelLabel} + + ); +}; + +export const MessageResponseDetailsSheet: FC<{ + open: boolean; + onOpenChange: (open: boolean) => void; +}> = ({ open, onOpenChange }) => { + const timing = useMessageTiming(); + const { + message, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + } = useResponseModelDisplay(); + const promptTokens = + usage?.promptTokens ?? asNumber(serverTimings?.prompt_n); + const completionTokens = + usage?.completionTokens ?? + timing?.tokenCount ?? + asNumber(serverTimings?.predicted_n); + const totalTokens = + usage?.totalTokens ?? + (promptTokens != null && completionTokens != null + ? promptTokens + completionTokens + : undefined); + const totalTime = + responseDetails?.durationMs ?? timing?.totalStreamTime ?? undefined; + const summaryLabel = + modelLabel === "Not recorded" ? "Model not recorded" : `Used ${modelLabel}`; + const messageToolCalls = toolCallsFromContent(message.content); + const toolCalls = + responseDetails?.toolCalls && responseDetails.toolCalls.length > 0 + ? responseDetails.toolCalls + : messageToolCalls; + + return ( + + + + + + Response details + + + Timing, model, token, and tool details for this response. + + + +
+
+

+ {summaryLabel} +

+ {providerLabel ? ( +

+ {providerLabel} +

+ ) : null} +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+ + + ); +}; diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index c5b53577cb..96d21d6fe7 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -6,6 +6,7 @@ /* eslint-disable react-refresh/only-export-components */ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { MessageResponseModelBadge } from "@/components/assistant-ui/message-response-details-sheet"; import { Collapsible, CollapsibleContent, @@ -390,14 +391,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ onOpenChange={handleOpenChange} variant={variant} > -
+
-
+ + + +
{isOpen && !isReasoningStreaming && ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 7fe98b701f..09551cd413 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -12,6 +12,10 @@ import { import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts"; +import { + MessageResponseDetailsSheet, + MessageResponseModelBadge, +} from "@/components/assistant-ui/message-response-details-sheet"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources"; @@ -3564,6 +3568,9 @@ const AssistantMessage: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); const messageContent = useAuiState(({ message }) => message.content); + const hasReasoningParts = useAuiState(({ message }) => + message.parts.some((part) => part.type === "reasoning"), + ); const incognito = useChatRuntimeStore((s) => s.incognito); // Use global store for editing state to ensure a single source of truth @@ -3620,7 +3627,7 @@ const AssistantMessage: FC = () => { return (
@@ -3649,6 +3656,11 @@ const AssistantMessage: FC = () => {
) : ( <> + {!hasReasoningParts ? ( +
+ +
+ ) : null} @@ -3893,58 +3905,76 @@ const EditAssistantMessageButton: FC = () => { const AssistantActionBar: FC = () => { const { forkMessage, forkDisabled } = useForkMessageAction(); + const [detailsOpen, setDetailsOpen] = useState(false); return ( - - - - - - - - - - - - - - + <> + + + + + + - - e.preventDefault()} - className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" - > - void forkMessage()} - className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + + + + + + + + + + e.preventDefault()} + className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" > - - Fork in new chat - - - + void forkMessage()} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + > + + Fork in new chat + + + + + Export as Markdown + + + setDetailsOpen(true)} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" + > - Export as Markdown + See response details - - - - - + + + + + + ); }; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index df266fd749..2ca6ceddf7 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -140,6 +140,32 @@ interface ServerTimings { diffusion_steps_per_second?: number; } +interface ResponseDetailsMetadata { + modelId: string; + modelLabel: string; + responseModelId: string; + providerId?: string; + providerName: string; + providerType: string; + startedAt: number; + finishedAt: number; + durationMs: number; + sessionId?: string; + cancelId: string; + toolCalls: string[]; + tools: { + search: boolean; + fetch: boolean; + code: boolean; + images: boolean; + mcp: boolean; + docs: boolean; + artifacts: boolean; + confirmToolCalls: boolean; + bypassPermissions: boolean; + }; +} + type RunMessages = Parameters[0]["messages"]; type RunMessage = RunMessages[number]; @@ -1769,6 +1795,9 @@ export function createOpenAIStreamAdapter( (provider) => provider.id === externalSelection.providerId, ) : null; + const selectedModelSummary = runtime.models.find( + (model) => model.id === params.checkpoint, + ); const externalApiKey = externalProvider ? getExternalProviderApiKey(externalProvider.id).trim() : ""; @@ -2151,6 +2180,7 @@ export function createOpenAIStreamAdapter( let waitingFirstChunk = true; let firstTokenSettled = false; const streamStartTime = Date.now(); + let responseModelId = externalSelection?.modelId ?? params.checkpoint; let firstTokenTime: number | undefined; let totalChunks = 0; let resolveFirstToken: (() => void) | null = null; @@ -2372,6 +2402,59 @@ export function createOpenAIStreamAdapter( const externalBackendProviderType = toExternalBackendProviderType( externalProvider?.providerType, ); + const buildResponseDetails = ( + finishedAt: number, + ): ResponseDetailsMetadata => ({ + modelId: params.checkpoint, + modelLabel: + (isExternalRequest || responseModelId !== params.checkpoint + ? responseModelId + : selectedModelSummary?.name || responseModelId) || + params.checkpoint || + "Unknown model", + responseModelId: + responseModelId || + externalSelection?.modelId || + params.checkpoint, + ...(externalProvider?.id ? { providerId: externalProvider.id } : {}), + providerName: + externalProvider?.name ?? + (isExternalRequest ? "External provider" : "Local model"), + providerType: externalProvider?.providerType ?? "local", + startedAt: streamStartTime, + finishedAt, + durationMs: finishedAt - streamStartTime, + ...(sandboxSessionId ? { sessionId: sandboxSessionId } : {}), + cancelId, + toolCalls: Array.from( + new Set( + toolCallParts + .map((part) => part.toolName) + .filter( + (toolName): toolName is string => + typeof toolName === "string" && toolName.length > 0, + ), + ), + ), + tools: { + search: + webSearchEnabledForThisTurn || + (!isExternalRequest && supportsTools && toolsEnabled), + fetch: webFetchEnabledForThisTurn, + code: + codeExecEnabledForThisTurn || + (!isExternalRequest && supportsTools && codeToolsEnabled), + images: imageGenerationEnabledForThisTurn, + mcp: !isExternalRequest && supportsTools && mcpEnabledForChat, + docs: + !isExternalRequest && + supportsTools && + (ragEnabled || projectRagEnabled), + artifacts: renderHtmlToolEnabledForThisTurn, + confirmToolCalls, + bypassPermissions, + }, + }); const externalCapabilities = getProviderCapabilities( externalProvider?.providerType, ); @@ -2768,6 +2851,11 @@ export function createOpenAIStreamAdapter( const stream = streamChatCompletions(requestPayload, abortSignal); for await (const chunk of stream) { + const chunkModel = (chunk as { model?: unknown }).model; + if (typeof chunkModel === "string" && chunkModel.length > 0) { + responseModelId = chunkModel; + } + // Handle tool status events const toolStatusText = ( chunk as unknown as { _toolStatus?: string } @@ -3435,11 +3523,12 @@ export function createOpenAIStreamAdapter( }); } + const finishedAt = Date.now(); const finalTiming = buildTiming( streamStartTime, totalChunks, serverPromptEvalTime ?? firstTokenTime, - Date.now() - streamStartTime, + finishedAt - streamStartTime, finalTokenCount, toolCallParts.length, finalTokPerSec, @@ -3475,6 +3564,7 @@ export function createOpenAIStreamAdapter( modelId: params.checkpoint, } : undefined, + responseDetails: buildResponseDetails(finishedAt), timing: finalTiming, }, }, diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 75749b0040..7cd9611c71 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -26,7 +26,12 @@ export { type PlusMenuItemId, } from "./stores/plus-menu-prefs-store"; export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; -export { isExternalModelId } from "./external-providers"; +export { + customProviderDisplayName, + isExternalModelId, + parseExternalModelId, +} from "./external-providers"; +export { useExternalProvidersStore } from "./stores/external-providers-store"; export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; diff --git a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts index f019b7dc3d..5b5011a78f 100644 --- a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts @@ -7,11 +7,14 @@ import { persist } from "zustand/middleware"; // Client-side chat UI prefs kept in localStorage, not the chat DB. // confirmDeleteChats: when off, deleting a chat skips the confirm dialog. // showModelDisclaimer: when off, hide the "LLMs can make mistakes" footer note. +// showResponseModel: when on, assistant responses show the producing model. export interface ChatPreferencesState { confirmDeleteChats: boolean; setConfirmDeleteChats: (value: boolean) => void; showModelDisclaimer: boolean; setShowModelDisclaimer: (value: boolean) => void; + showResponseModel: boolean; + setShowResponseModel: (value: boolean) => void; } export const useChatPreferencesStore = create()( @@ -23,6 +26,9 @@ export const useChatPreferencesStore = create()( showModelDisclaimer: true, setShowModelDisclaimer: (showModelDisclaimer) => set({ showModelDisclaimer }), + showResponseModel: false, + setShowResponseModel: (showResponseModel) => + set({ showResponseModel }), }), { name: "unsloth_chat_preferences", @@ -32,6 +38,7 @@ export const useChatPreferencesStore = create()( ...current, confirmDeleteChats: saved?.confirmDeleteChats ?? true, showModelDisclaimer: saved?.showModelDisclaimer ?? true, + showResponseModel: saved?.showResponseModel ?? false, }; }, }, diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx index b26b09e022..9519898b21 100644 --- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx @@ -213,6 +213,12 @@ export function ChatTab() { const setShowModelDisclaimer = useChatPreferencesStore( (state) => state.setShowModelDisclaimer, ); + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const setShowResponseModel = useChatPreferencesStore( + (state) => state.setShowResponseModel, + ); useEffect(() => { void countAllChats().then(setCount); @@ -412,6 +418,15 @@ export function ChatTab() { onCheckedChange={setShowModelDisclaimer} /> + + + diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py new file mode 100644 index 0000000000..04301a0de6 --- /dev/null +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -0,0 +1,93 @@ +"""Static contract for the chat response-details action and metadata.""" + +from __future__ import annotations + +import re +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +THREAD_TSX = REPO / "studio/frontend/src/components/assistant-ui/thread.tsx" +DETAILS_TSX = ( + REPO / "studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx" +) +REASONING_TSX = REPO / "studio/frontend/src/components/assistant-ui/reasoning.tsx" +ADAPTER_TS = REPO / "studio/frontend/src/features/chat/api/chat-adapter.ts" +CHAT_PREFS_TS = REPO / "studio/frontend/src/features/chat/stores/chat-preferences-store.ts" +CHAT_TAB_TSX = REPO / "studio/frontend/src/features/settings/tabs/chat-tab.tsx" + + +def test_assistant_more_menu_exposes_response_details_action(): + src = THREAD_TSX.read_text() + assert "MessageResponseDetailsSheet" in src + assert "See response details" in src + assert "setDetailsOpen(true)" in src + + +def test_response_details_sheet_uses_unsloth_sheet_and_key_sections(): + src = DETAILS_TSX.read_text() + assert "SheetContent" in src + assert "Response details" in src + assert "MessageResponseModelBadge" in src + assert "showResponseModel" in src + assert "ChipIcon" not in src + assert "s.params.checkpoint" not in src + assert "Not recorded" in src + assert "min-w-0 break-words font-heading" in src + assert "toolCallsFromContent(message.content)" in src + assert 'label="Called"' in src + for section in ["Response", "Tokens", "Timing", "Tools"]: + assert f'title="{section}"' in src + for field in ["Model", "Provider", "Total", "Cache hits", "Enabled", "Called"]: + assert f'label="{field}"' in src + + +def test_response_model_chip_is_user_configurable_and_rendered_in_metadata_rows(): + prefs_src = CHAT_PREFS_TS.read_text() + chat_tab_src = CHAT_TAB_TSX.read_text() + thread_src = THREAD_TSX.read_text() + reasoning_src = REASONING_TSX.read_text() + + assert "showResponseModel: boolean" in prefs_src + assert "showResponseModel: false" in prefs_src + assert "showResponseModel: saved?.showResponseModel ?? false" in prefs_src + assert "Show response model" in chat_tab_src + assert "setShowResponseModel" in chat_tab_src + assert "aui-response-model-badge inline-flex min-h-5" in DETAILS_TSX.read_text() + assert "leading-5" in DETAILS_TSX.read_text() + assert "group-hover/assistant-message:opacity-100" in DETAILS_TSX.read_text() + assert "MessageResponseModelBadge" in thread_src + assert "hasReasoningParts" in thread_src + assert "group/assistant-message aui-assistant-message-root" in thread_src + assert "pointer-events-none relative h-0" in thread_src + assert "MessageResponseModelBadge" in reasoning_src + assert 'className="min-w-0 flex-none"' in reasoning_src + assert "hidden min-w-0 max-w-[12rem]" in reasoning_src + assert "group-hover/assistant-message:inline-flex" in reasoning_src + + +def test_response_details_metadata_is_persisted_without_backend_schema_change(): + src = ADAPTER_TS.read_text() + assert "interface ResponseDetailsMetadata" in src + assert "buildResponseDetails" in src + assert "responseDetails: buildResponseDetails(finishedAt)" in src + assert "toolCalls: Array.from(" in src + assert "!isExternalRequest && supportsTools && toolsEnabled" in src + assert "!isExternalRequest && supportsTools && codeToolsEnabled" in src + assert re.search(r"selectedModelSummary\?\.name\s*\|\|\s*responseModelId", src) + assert "providerName" in src + assert "cancelId" in src + metadata_block = src[ + src.find("interface ResponseDetailsMetadata") : src.find("type RunMessages") + ] + builder_block = src[ + src.find("const buildResponseDetails") : src.find("const externalCapabilities") + ] + for forbidden in [ + "encrypted_api_key", + "externalApiKey", + "apiKey", + "providerKey", + "secret", + ]: + assert forbidden not in metadata_block + assert forbidden not in builder_block From 8efcc17f476c21dd6f5534ce719b014d7e6a1e4c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 07:20:31 -0700 Subject: [PATCH 12/19] Studio: account for DeepSeek-V4 compute buffer in context auto-fit (#6940) * Studio: account for DeepSeek-V4 compute buffer in context auto-fit DeepSeek-V4-Flash's lightning indexer plus compressed sparse attention reserve a large context-scaling compute buffer that _compute_buffer_ctx_bytes did not model (the KQ-mask and dequant-scratch rates both miss it, even with an f16 cache). Measured on UD-Q4_K_XL at ub 512 it is about 65.5 GiB at 1M context, which the mask estimate puts near 1.5 GiB, so the auto-fit kept the full 1M train context and llama-server OOM'd allocating the ~70 GB buffer, then spilled to CPU (~4 tok/s). Add a deepseek4-gated flat plus per-token term so the fit caps the context (about 256k on a B200) and the model stays fully on GPU. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 14 +++++ studio/backend/tests/test_compute_buffer.py | 61 +++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8b40f5fccd..11a4ebb3ec 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3143,6 +3143,12 @@ class LlamaCppBackend: _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) + # DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large + # context-scaling compute buffer the rates above miss (present even with an f16 + # cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without + # it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU. + _DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch + _DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M) def _estimate_compute_buffer_bytes( self, @@ -3192,6 +3198,14 @@ class LlamaCppBackend: if n_embd <= 0 or n_ctx <= 0: return 0 ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if getattr(self, "_architecture", None) == "deepseek4": + # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires + # for any KV type -- the indexer scratch is present even with an f16 cache. + ub_scale = ub / self._DEFAULT_N_UBATCH + return int( + self._DSV4_CTX_COMPUTE_FLAT_BYTES + + self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale + ) if _kv_bytes_per_elem(cache_type_kv) < 2.0: # Quantized cache: the dequant scratch dominates and scales with n_embd. # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 5d14c5c5bd..8408f8203d 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -65,12 +65,14 @@ def _backend( vocab = 248320, embd = 5120, mla = None, + arch = None, ): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd b._key_length_mla = mla # non-None -> MLA (compressed attention) + b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4') return b @@ -290,3 +292,62 @@ class TestContextBufferMLA: b = _backend(embd = 6144, mla = 256) est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB assert est <= 4141 * 1.7 + + +class TestContextBufferDSV4: + """DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention + compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache). + Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit + must see this so it does not commit the full 1M train context and OOM (spilling + to CPU at ~4 tok/s).""" + + _MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx + GIB = 1024**3 + + def test_covers_measured_1m_buffer(self): + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}" + + def test_not_wildly_over_at_1m(self): + # Within ~1.3x of measured so the fit still grants a large (~256k) context. + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib <= self._MEASURED_1M_GIB * 1.3 + + def test_fires_for_f16_cache(self): + # The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must + # reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx. + dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + assert dsv4 > 40 * other + + def test_cache_type_independent(self): + # Indexer scratch is present for an f16 and a quantized cache alike. + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + + def test_flat_floor_at_small_ctx(self): + # ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB). + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0 + + def test_scales_with_context_and_ubatch(self): + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536) + assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes( + 131072, n_ubatch = 256 + ) + + def test_non_dsv4_unchanged(self): + # Regression guard: a non-deepseek4 model keeps the mask-only f16 rate. + b = _backend(embd = 4096, arch = "llama") + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000 + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY + assert per_tok == pytest.approx(expected, rel = 1e-6) From 37075c542258e87bb556f6ed7496ea88a818c3b6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 07:49:59 -0700 Subject: [PATCH 13/19] Bump install.sh / install.ps1 pin to unsloth>=2026.7.1 (#6943) Co-authored-by: danielhanchen --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 8c667df079..9114f80af9 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2155,7 +2155,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2169,7 +2169,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2235,7 +2235,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2247,7 +2247,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2275,7 +2275,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index 14fbba478d..796d80e401 100755 --- a/install.sh +++ b/install.sh @@ -2706,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2721,7 +2721,7 @@ if [ "$_MIGRATED" = true ]; then # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2925,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2943,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2975,7 +2975,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From 07ecdb34c092cb0107f74dbf117616e320ec6ff2 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:24:32 +0530 Subject: [PATCH 14/19] Sort chat recents by last activity (#6844) * show chat by by last activity * Update chat thread updated_at logic and enhance sidebar chat item handling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/chat_history.py | 4 +- studio/backend/storage/studio_db.py | 77 +++++++++- .../tests/test_chat_history_storage.py | 134 ++++++++++++++++++ .../frontend/src/components/app-sidebar.tsx | 16 ++- .../chat/hooks/use-chat-sidebar-items.ts | 26 +++- studio/frontend/src/features/chat/types.ts | 1 + .../chat/utils/chat-history-storage.ts | 5 +- studio/frontend/src/i18n/locales/en.ts | 1 + studio/frontend/src/i18n/locales/zh-CN.ts | 1 + 9 files changed, 252 insertions(+), 13 deletions(-) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 963d584303..7a27a58a52 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -56,6 +56,7 @@ class ChatThread(BaseModel): projectId: Optional[str] = None archived: bool = False createdAt: int + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None forkedFromThreadId: Optional[str] = None @@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel): projectId: Optional[str] = None archived: Optional[bool] = None createdAt: Optional[int] = None + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None @@ -252,7 +254,7 @@ async def patch_thread( current_subject: str = Depends(get_current_subject), ): patch = payload.model_dump(exclude_unset = True) - for field in ("title", "modelType", "modelId", "archived", "createdAt"): + for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): if field in patch and patch[field] is None: raise HTTPException(status_code = 400, detail = f"{field} cannot be null") if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 41a9adcc29..87aa50ee26 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -240,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: project_id TEXT, archived INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, + updated_at INTEGER, openai_code_exec_container_id TEXT, anthropic_code_exec_container_id TEXT, forked_from_thread_id TEXT, @@ -261,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") if "forked_from_message_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT") + if "updated_at" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER") + # Floor at created_at: forked threads copy older ancestor messages, + # so the fork's creation time must win over the branch message times. + conn.execute( + """ + UPDATE chat_threads SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + """ + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -992,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "projectId": data.get("project_id") or None, "archived": bool(data["archived"]), "createdAt": data["created_at"], + "updatedAt": data.get("updated_at") + if data.get("updated_at") is not None + else data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), "forkedFromThreadId": data.get("forked_from_thread_id"), @@ -1039,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, @@ -1049,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict: project_id = excluded.project_id, archived = excluded.archived, created_at = excluded.created_at, + updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at), openai_code_exec_container_id = excluded.openai_code_exec_container_id, anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id, forked_from_thread_id = excluded.forked_from_thread_id, @@ -1063,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict: thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), + int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None, thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), thread.get("forkedFromThreadId"), @@ -1084,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "projectId": ("project_id", patch.get("projectId")), "archived": ("archived", 1 if patch.get("archived") else 0), "createdAt": ("created_at", patch.get("createdAt")), + "updatedAt": ("updated_at", patch.get("updatedAt")), "openaiCodeExecContainerId": ( "openai_code_exec_container_id", patch.get("openaiCodeExecContainerId"), @@ -1155,7 +1180,8 @@ def list_chat_threads( conn = get_connection() try: rows = conn.execute( - f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC", + f"SELECT * FROM chat_threads {where} " + "ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC", values, ).fetchall() return [_chat_thread_from_row(row) for row in rows] @@ -1394,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts( ) +def _bump_chat_thread_updated_at( + conn: sqlite3.Connection, thread_id: str, message_created_at: int +) -> None: + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX(COALESCE(updated_at, created_at), ?) + WHERE id = ? + """, + (message_created_at, thread_id), + ) + + +def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None: + """Set updated_at from the remaining messages, floored at created_at. + + Unlike the ratchet-only bump, this can lower updated_at -- needed after + pruning, which may delete the thread's newest message. + """ + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + WHERE id = ? + """, + (thread_id,), + ) + + def upsert_chat_message(message: dict) -> dict: conn = get_connection() try: @@ -1432,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict: int(message["createdAt"]), ), ) + _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) conn.commit() return message except Exception: @@ -1484,6 +1549,12 @@ def sync_chat_messages( for m in messages ], ) + if prune_missing: + _recompute_chat_thread_updated_at(conn, thread_id) + elif messages: + _bump_chat_thread_updated_at( + conn, thread_id, max(int(m["createdAt"]) for m in messages) + ) conn.commit() return list_chat_messages(thread_id) except ChatMessageConflictError: diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index aa19df15fe..0239410734 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -4,6 +4,7 @@ import os import platform import shutil +import sqlite3 import threading import uuid from pathlib import Path @@ -11,6 +12,7 @@ from pathlib import Path import pytest from storage import studio_db +from utils.paths import studio_db_path def _reset_studio_db( @@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] +def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + assert thread["updatedAt"] == thread["createdAt"] + + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 1_700_000_001_000, "newer")], + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + +def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1_700_000_000_500, "older"), + _message("msg-2", 1_700_000_001_000, "newest"), + ], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + # Pruning the newest message must lower updated_at to the remaining one. + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-1", 1_700_000_000_500, "older")], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + # Pruning every message falls back to created_at. + studio_db.sync_chat_messages("thread-1", [], prune_missing = True) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"] + + +def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + + studio_db.upsert_chat_thread(_thread()) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + +def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + older = _thread("thread-old") + older["createdAt"] = 1_700_000_000_000 + newer = _thread("thread-new") + newer["createdAt"] = 1_700_000_100_000 + studio_db.upsert_chat_thread(older) + studio_db.upsert_chat_thread(newer) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"] + + studio_db.upsert_chat_message( + _message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old") + ) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"] + + +def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + db_path = studio_db_path() + db_path.parent.mkdir(parents = True, exist_ok = True) + conn = sqlite3.connect(str(db_path)) + try: + conn.execute( + """ + CREATE TABLE chat_threads ( + id TEXT NOT NULL PRIMARY KEY, + title TEXT NOT NULL, + model_type TEXT NOT NULL, + model_id TEXT, + pair_id TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE chat_messages ( + id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL, + parent_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + attachments_json TEXT, + metadata_json TEXT, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-with-msgs", "Old", "base", 1_700_000_000_000), + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-empty", "Empty", "base", 1_700_000_050_000), + ) + # Fork-like thread: copied ancestor messages predate the thread itself. + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-fork", "Fork", "base", 1_700_000_100_000), + ) + conn.executemany( + "INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)", + [ + ("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000), + ("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000), + ("m3", "thread-fork", "user", "[]", 1_700_000_001_000), + ], + ) + conn.commit() + finally: + conn.close() + + assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000 + assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000 + + def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) project = studio_db.upsert_chat_project(_project()) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index fb1a1fc9c7..2dd0d02515 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -359,7 +359,11 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ + const { + items: allChatItems, + archivedItems: archivedChatItems, + loaded: chatItemsLoaded, + } = useChatSidebarItems({ enabled: !isStudioRoute, requireMessages: false, }); @@ -1306,6 +1310,16 @@ export function AppSidebar() { renderChatSidebarItem(item, "recent"), )} + {/* "No chats yet" only when there is truly no history: + project-scoped and archived threads leave Recents empty + but still count as existing chats. */} + {chatItemsLoaded && + allChatItems.length === 0 && + archivedChatItems.length === 0 && ( +

+ {t("shell.navigation.noChatsYet")} +

+ )} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index 55d7777e43..0a0df1139b 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -27,16 +27,21 @@ export interface SidebarItem { id: string; title: string; createdAt: number; + updatedAt: number; isFork?: boolean; projectId?: string | null; } +function lastActivityAt(thread: ThreadRecord): number { + return thread.updatedAt ?? thread.createdAt; +} + export function groupThreads( threads: ThreadRecord[], archived = false, ): SidebarItem[] { const items: SidebarItem[] = []; - const seenPairs = new Set(); + const pairItems = new Map(); for (const t of threads) { // Coerce archived to a boolean before comparing. Legacy threads (from the @@ -48,30 +53,35 @@ export function groupThreads( continue; } if (t.pairId) { - if (seenPairs.has(t.pairId)) { + const existing = pairItems.get(t.pairId); + if (existing) { + existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t)); continue; } - seenPairs.add(t.pairId); - items.push({ + const item: SidebarItem = { type: "compare", id: t.pairId, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), projectId: t.projectId ?? null, - }); + }; + pairItems.set(t.pairId, item); + items.push(item); } else if (!t.pairId) { items.push({ type: "single", id: t.id, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), isFork: Boolean(t.forkedFromThreadId), projectId: t.projectId ?? null, }); } } - return items.sort((a, b) => b.createdAt - a.createdAt); + return items.sort((a, b) => b.updatedAt - a.updatedAt); } // Streaming fires CHAT_HISTORY_UPDATED_EVENT per chunk. Debounce so each quiet @@ -84,6 +94,7 @@ export function useChatSidebarItems(options?: { requireMessages?: boolean; }) { const [allThreads, setAllThreads] = useState([]); + const [loaded, setLoaded] = useState(false); const enabled = options?.enabled ?? true; const requireMessages = options?.requireMessages ?? true; @@ -111,6 +122,7 @@ export function useChatSidebarItems(options?: { // were in flight, or if the effect was torn down. if (cancelled || seq !== requestSeq) return; setAllThreads(threads); + setLoaded(true); } catch (error) { if (isExpectedBackgroundChatStorageError(error)) { return; @@ -144,7 +156,7 @@ export function useChatSidebarItems(options?: { const archivedItems = groupThreads(allThreads ?? [], true); const canCompare = useChatRuntimeStore((s) => Boolean(s.params.checkpoint)); - return { items, archivedItems, canCompare }; + return { items, archivedItems, canCompare, loaded }; } function cancelIfRunning(threadId: string): void { diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 3fe69ccc26..111510d925 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -36,6 +36,7 @@ export interface ThreadRecord { projectId?: string | null; archived: boolean; createdAt: number; + updatedAt?: number; /** * OpenAI shell tool container id from a prior response. When set, the * next turn reuses it via `environment.type="container_reference"` so diff --git a/studio/frontend/src/features/chat/utils/chat-history-storage.ts b/studio/frontend/src/features/chat/utils/chat-history-storage.ts index 4294887df0..00df3657b1 100644 --- a/studio/frontend/src/features/chat/utils/chat-history-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-history-storage.ts @@ -590,7 +590,10 @@ export async function listStoredChatThreads( } return Array.from(byId.values()) .filter((thread) => matchesThreadListArgs(thread, args)) - .sort((a, b) => b.createdAt - a.createdAt); + .sort( + (a, b) => + (b.updatedAt ?? b.createdAt) - (a.updatedAt ?? a.createdAt), + ); } export async function listStoredChatThreadsWithMessages( diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 3d73ad4343..b67fd5ca1d 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -41,6 +41,7 @@ export const en = { recipes: "Recipes", export: "Export", recents: "Recents", + noChatsYet: "No chats yet", settings: "Settings", api: "API", lightMode: "Light Mode", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 5fd31dc7cb..f6dc265fc7 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -41,6 +41,7 @@ export const zhCN = { recipes: "配方", export: "导出", recents: "最近", + noChatsYet: "暂无对话", settings: "设置", api: "API", lightMode: "浅色模式", From 93c9d6d0dd0bfd48d5766ac952c3a880a47b3727 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 15:13:53 -0300 Subject: [PATCH 15/19] Studio: render \[ \] and \( \) LaTeX delimiters in chat (#6914) --- studio/frontend/src/lib/latex.ts | 169 ++++++++++++++++++++++++++++--- 1 file changed, 155 insertions(+), 14 deletions(-) diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 1a4ebdace8..86a9634048 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -1,8 +1,13 @@ // Adapted from LibreChat's latex.ts // https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts // -// Escapes currency dollar signs so they are not misinterpreted as LaTeX math -// delimiters when singleDollarTextMath is enabled. +// Two jobs, in order: +// 1. Convert LaTeX bracket delimiters (`\[...\]`, `\(...\)`) into the dollar +// forms remark-math understands (`$$...$$`, `$...$`). remark-math only +// tokenizes dollar delimiters, so models that emit `\[...\]` / `\(...\)` +// would otherwise render as literal text. +// 2. Escape currency dollar signs so they are not misinterpreted as LaTeX +// math delimiters when singleDollarTextMath is enabled. /** * Matches a single $ followed by a number pattern (currency), e.g.: @@ -15,14 +20,14 @@ const CURRENCY_REGEX = /(? { const regions: Array<[number, number]> = []; - // Fenced code blocks: ```...``` - const fencedRe = /```[\s\S]*?```/g; + // Fenced code blocks: ```...``` and ~~~...~~~ (both are code in GFM) + const fencedRe = /```[\s\S]*?```|~~~[\s\S]*?~~~/g; let match: RegExpExecArray | null; while ((match = fencedRe.exec(content)) !== null) { regions.push([match.index, match.index + match[0].length]); @@ -51,9 +56,38 @@ function findCodeBlockRegions(content: string): Array<[number, number]> { } /** - * Binary search to check if a position falls inside any code region. + * Match an inline link/image `[text](DEST)`, capturing the destination as group 1 + * with the `d` flag so its span is read straight from `match.indices` (the text + * can contain an escaped `\](`, so a string search for the separator is unsafe). + * The text disallows unescaped `]`; the destination allows escapes and one level + * of balanced parens. */ -function isInCodeBlock( +const LINK_DEST_RE = + /!?\[(?:\\.|[^\]\\])*?\]\(((?:\\.|[^()\\]|\([^()]*\))*)\)/gd; + +/** + * Find the destination spans of inline links/images, so a `\(...\)` written with + * escaped parens inside a URL isn't rewritten as math (which would break the + * link). Only the destination is returned, not the link text, so math in the + * visible text still converts. Sorted, non-overlapping (matches are disjoint). + */ +function findLinkDestinationRegions(content: string): Array<[number, number]> { + if (!content.includes("](")) return []; + const regions: Array<[number, number]> = []; + let match: RegExpExecArray | null; + LINK_DEST_RE.lastIndex = 0; + while ((match = LINK_DEST_RE.exec(content)) !== null) { + // `indices` is present (the `d` flag); group 1 spans the destination. + regions.push(match.indices![1]); + } + return regions; +} + +/** + * Binary search to check if a position falls inside any region. Regions must be + * sorted by start and non-overlapping. + */ +function isInRegion( position: number, regions: Array<[number, number]>, ): boolean { @@ -174,9 +208,109 @@ function hasInlineMathCloser(content: string, offset: number): boolean { } /** - * Preprocess a markdown string to escape currency dollar signs so they are not - * parsed as LaTeX math delimiters. + * Matches a `\[...\]` (display) or `\(...\)` (inline) LaTeX span. Non-greedy so + * the first closer wins; dotall so display spans can wrap lines. `(? block `$$...$$` and `\(...\)` -> inline `$...$` so + * remark-math can tokenize them. Bodies are trimmed: remark-math won't open an + * inline span on `$ ` (a `$` followed by whitespace), and display fences must + * sit on their own line to render as a centered block (not inline math), so + * `\[...\]` becomes `\n$$\n...\n$$\n`. * + * Spans inside code blocks/spans are left intact (a code sample showing `\(x\)` + * must not be rewritten). + * + * A space is inserted between a converted span and a following `$` so their + * delimiters can't fuse (`\(a\)\(b\)` -> `$a$$b$` would mis-tokenize into one + * broken span). A preceding currency (`$5\(x\)`) is instead broken later by the + * currency escape pass. + * + * Returns the rewritten text and the `[start, end)` ranges (in the rewritten + * string) of every span it produced, so the currency pass can skip them. + */ +function convertLatexDelimiters(content: string): { + text: string; + mathRegions: Array<[number, number]>; +} { + if (!content.includes("\\[") && !content.includes("\\(")) { + return { text: content, mathRegions: [] }; + } + + const codeRegions = findCodeBlockRegions(content); + const linkRegions = findLinkDestinationRegions(content); + const inSkipZone = (pos: number) => + isInRegion(pos, codeRegions) || isInRegion(pos, linkRegions); + // Pushed in ascending, non-overlapping order (offset only grows), so this + // stays valid for isInRegion's binary search without a sort. + const mathRegions: Array<[number, number]> = []; + // Accumulate into an array, not a string: reading the last char off a growing + // `+=` accumulator flattens its rope every append (O(n^2) over many spans, on + // the per-frame streaming path), so track the tail char and length instead. + const parts: string[] = []; + let offset = 0; + let lastChar = ""; + let last = 0; + // Append a chunk, separating a trailing `$` from a leading `$` so two spans + // can't fuse. Returns where the chunk landed (after any inserted space). + const append = (chunk: string): number => { + if (!chunk) return offset; + if (lastChar === "$" && chunk.startsWith("$")) { + parts.push(" "); + offset += 1; + } + const start = offset; + parts.push(chunk); + offset += chunk.length; + lastChar = chunk[chunk.length - 1]; + return start; + }; + let match: RegExpExecArray | null; + LATEX_DELIM_RE.lastIndex = 0; + while ((match = LATEX_DELIM_RE.exec(content)) !== null) { + const matchEnd = match.index + match[0].length; + // Skip if either delimiter is inside code or a link destination: an opener + // outside such a zone must not consume a closer inside one and rewrite + // across the boundary. Resume right after this opener (not past the whole + // match) so a valid span that this match spanned across (a stray code `\(` + // paired with a real closer) is still found on the next pass, not swallowed. + if (inSkipZone(match.index) || inSkipZone(matchEnd - 1)) { + LATEX_DELIM_RE.lastIndex = match.index + 1; + continue; + } + const isDisplay = match[1] !== undefined; + const body = (isDisplay ? match[1] : match[2]).trim(); + // Leave an empty span (`\(\)`) literal; a bare `$$` would open a stray + // display block that swallows following text. + if (!body) { + continue; + } + append(content.slice(last, match.index)); + const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`; + const start = append(wrapped); + mathRegions.push([start, offset]); + last = matchEnd; + } + append(content.slice(last)); + return { text: parts.join(""), mathRegions }; +} + +/** + * Preprocess a markdown string so LaTeX renders: convert bracket delimiters to + * dollar forms, then escape currency dollar signs so they are not parsed as + * math delimiters. + * + * - `\[E = mc^2\]` becomes a `$$` display block on its own lines (display math) + * - `\(\alpha\)` becomes `$\alpha$` (inline math) + * - `\(x\)` in a code span is untouched * - `$5` alone becomes `\$5` (currency, not math) * - `$\alpha$` is untouched (real LaTeX) * - `$30^\circ$` is untouched (LaTeX whose body starts with a digit) @@ -185,15 +319,22 @@ function hasInlineMathCloser(content: string, offset: number): boolean { * - Currency inside code blocks/spans is untouched */ export function preprocessLaTeX(content: string): string { - if (!content.includes("$")) return content; + const { text, mathRegions } = convertLatexDelimiters(content); - const codeRegions = findCodeBlockRegions(content); + if (!text.includes("$")) return text; - return content.replace(CURRENCY_REGEX, (match, offset) => { - if (isInCodeBlock(offset, codeRegions)) { + const codeRegions = findCodeBlockRegions(text); + + return text.replace(CURRENCY_REGEX, (match, offset) => { + if (isInRegion(offset, codeRegions)) { return match; } - if (hasInlineMathCloser(content, offset)) { + // Skip the spans we just created from `\(...\)` so a numeric body like + // `$5$` isn't re-escaped back to literal `\$5$`. + if (isInRegion(offset, mathRegions)) { + return match; + } + if (hasInlineMathCloser(text, offset)) { return match; } return "\\" + match; From 304b8eca7ae8a7bd743850a248308931c062dab9 Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Tue, 7 Jul 2026 23:58:18 +0530 Subject: [PATCH 16/19] fix: match qwen3-thinking double-newline in train_on_responses_only response pattern (#6926) * fix: match qwen3-thinking chat template double-newline in response pattern The Qwen3-thinking chat template generates `\n\n` (double newline) after the think tag, but `train_on_responses_only` was looking for `\n` (single newline). `\n\n` is token 271 while `\n` is token 198 -- different tokens, so the pattern match in `train_on_responses_only` fails, masking ALL tokens and dropping 100% of training samples. Update the response pattern from `\n` to `\n\n` to match what the actual qwen3-thinking template generates. Fixes #6919 * fix qwen3 thinking response marker --------- Co-authored-by: Ayushman Paul Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- studio/backend/utils/datasets/model_mappings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 463d26a692..9d2c983aed 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -487,7 +487,7 @@ TEMPLATE_TO_RESPONSES_MAPPER = { }, "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", From a9db53e189f2c23586bfc4a6f472448afc4807ef Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 19:50:40 -0300 Subject: [PATCH 17/19] Studio: stream reasoning tokens in the tool-loop generator (fixes DeepSeek thinking not streaming with a pill on) (#6947) --- .../core/inference/anthropic_compat.py | 28 +++ studio/backend/core/inference/llama_cpp.py | 54 +++- .../backend/tests/test_anthropic_messages.py | 43 ++++ .../backend/tests/test_llama_cpp_tool_loop.py | 231 +++++++++++++++++- 4 files changed, 338 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 7b572a28ff..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -258,6 +258,10 @@ class AnthropicStreamEmitter: self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" + # Net minus in the text emitted to the client. Tracked + # from emitted deltas (not _prev_text, which a final bare shrink clobbers) + # so an unclosed reasoning-only block can be balanced before close. + self._open_think_tags: int = 0 self._usage: dict = {} def start( @@ -317,6 +321,7 @@ class AnthropicStreamEmitter: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: + events.extend(self._close_open_think()) events.append(self._close_block()) self._open_tool_call_id = None self._open_tool_use_id = None @@ -344,12 +349,33 @@ class AnthropicStreamEmitter: ) return events + def _close_open_think(self) -> list[str]: + """Emit a ```` delta when the streamed text left a ```` + open. This emitter diffs cumulative snapshots and drops the generator's + final bare shrink, so a reasoning-only reply would otherwise end on an + unclosed tag. Mirrors the chat route's reasoning extractor, which closes + the block on finish; balances the block before it is closed.""" + if not self._text_block_open or self._open_think_tags <= 0: + return [] + self._open_think_tags = 0 + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": ""}, + }, + ) + ] + def _handle_content(self, event: dict) -> list[str]: cumulative = event.get("text", "") new_text = cumulative[len(self._prev_text) :] self._prev_text = cumulative if not new_text: return [] + self._open_think_tags += new_text.count("") - new_text.count("") if not self._text_block_open: events = self._open_text_block() else: @@ -374,6 +400,7 @@ class AnthropicStreamEmitter: events = [] if self._text_block_open: + events.extend(self._close_open_think()) events.append(self._close_block()) # Defensive: close a stale open tool_use block before starting another. elif self._open_tool_call_id is not None: @@ -452,6 +479,7 @@ class AnthropicStreamEmitter: events.extend(self._open_text_block()) # Reset text tracking for the next synthesis turn self._prev_text = "" + self._open_think_tags = 0 return events def _open_text_block(self) -> list[str]: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 11a4ebb3ec..757467a008 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8603,13 +8603,31 @@ class LlamaCppBackend: } def _flush_reasoning_and_buffer(): - """Append buffered reasoning (as a block) then the held + """Close a live-streamed block (or emit the buffered reasoning + as one block if it never streamed), then append the held content_buffer to the cumulative display text.""" - nonlocal cumulative_display - if reasoning_accum: + nonlocal cumulative_display, in_thinking + if in_thinking: + cumulative_display += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed before a tool call drains, so + consumers without a reasoning extractor (Anthropic) get a balanced + block. Returns True when the caller should yield the result.""" + nonlocal cumulative_display, in_thinking, _last_emitted + if not in_thinking: + return False + cumulative_display += "" + in_thinking = False + if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + _last_emitted = cumulative_display + return True + return False + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" probe = strip_llama3_leading_sentinels(text.lstrip()) @@ -8797,6 +8815,10 @@ class LlamaCppBackend: # the structured tool call. has_structured_tc = True detect_state = _S_DRAINING + # Close the reasoning prefix before the tool card + # (mirrors the is_match path). + if _close_streamed_think(): + yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -8882,17 +8904,17 @@ class LlamaCppBackend: continue # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). + # Stream live except while DRAINING: reasoning is + # orthogonal to tool detection (content_buffer + # only), and the route resets prev_text on + # tool_start, so the block stays a + # monotonic prefix like the no-tool path. reasoning = delta.get("reasoning_content", "") if reasoning: if _reasoning_started_at is None: _reasoning_started_at = time.monotonic() reasoning_accum += reasoning - if detect_state == _S_STREAMING: + if detect_state != _S_DRAINING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -9020,9 +9042,15 @@ class LlamaCppBackend: _hold_buffer = True if _drain_silently: - # No visible prefix -- the buffered text IS - # the call; drain without yielding it. + # The buffered content IS the call; drain it + # without yielding. A live prefix is + # separate from it -- close that. detect_state = _S_DRAINING + if _close_streamed_think(): + yield { + "type": "content", + "text": cumulative_display, + } elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the @@ -9115,7 +9143,9 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only reply: show it as plain text. + # Reasoning-only reply: show it as the main response, + # not a thinking block (mirrors the no-tool path; the + # route's extractor closes the streamed ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 170b456eac..0c6550a3bb 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def _emitter_client_text(events: list[str]) -> str: + """Concatenate the text_delta payloads an SSE event list carries.""" + text = "" + for line in events: + for raw in line.split("\n"): + raw = raw.strip() + if not raw.startswith("data: "): + continue + data = json.loads(raw[len("data: ") :]) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + return text + + +def test_anthropic_emitter_closes_reasoning_only_think_block(): + # A reasoning-only reply streams X live then shrinks to bare X at EOF. + # This emitter diffs cumulative snapshots and drops the shrink, so without a + # closing pass the client text would end on an unclosed . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + # The generator's final bare-text shrink (dropped by the cumulative diff). + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fb1b0e52b7..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -221,7 +221,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" -def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): +def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): stream = [ _sse({"reasoning_content": "I am thinking."}), _sse({"reasoning_content": " Still thinking."}), @@ -240,17 +240,236 @@ def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): ) ) + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streams live during BUFFERING instead of arriving as one block: + # each reasoning delta is emitted immediately, wrapped in . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + summary_index = next( i for i, event in enumerate(events) if event["type"] == "reasoning_summary" ) - content_index = next(i for i, event in enumerate(events) if event["type"] == "content") - assert summary_index < content_index + final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + assert summary_index < final_content_index assert events[summary_index]["duration_ms"] == 62000 - assert ( - events[content_index]["text"] - == "I am thinking. Still thinking.Final answer." + + +def test_reasoning_streams_incrementally_with_tools(monkeypatch): + # Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the + # tool-loop generator must stream reasoning token-by-token like the no-tool + # path, not accumulate it and dump one buffered block. + stream = [ + _sse({"reasoning_content": "Step one."}), + _sse({"reasoning_content": " Step two."}), + _sse({"reasoning_content": " Step three."}), + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "think then answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) ) + reasoning_stage = [ + e["text"] + for e in events + if e["type"] == "content" + and e["text"].startswith("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.Done." + + +def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): + # A reasoning-only turn (whole answer in reasoning_content, no content, no + # tool) with a tool active streams the reasoning live, then resolves to the + # bare reasoning text -- identical to the no-tool generate_chat_completion + # path -- so the non-streaming drain still returns it as `content`, not an + # empty answer. + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 5.0, 5.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "just think"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streamed live during BUFFERING (the fix). + assert content_texts[0] == "The capital of France is Paris." + # Resolves to bare reasoning, matching the no-tool sibling. + assert content_texts[-1] == "The capital of France is Paris." + + +def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): + # Regression: reasoning streamed live during BUFFERING must be closed with + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . Mirrors the is_match (XML tool signal) path. + tool_stream = [ + _sse({"reasoning_content": "Let me search."}), + *_structured_tool_call("web_search", {"query": "weather"}, "call_1"), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + # Reasoning streamed live, then closed before the tool -- balanced block. + assert content_before_tool[0] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]: + """Replay the route's cumulative suffix-diff + reasoning extractor (the + shared core of routes/inference.py gguf_stream_chunks and the tool-loop + consumer) over content snapshots. Returns (visible, reasoning).""" + from routes.inference import _ResponsesReasoningExtractor + + extractor = _ResponsesReasoningExtractor(parse_think_markers = True) + prev_text = "" + visible: list[str] = [] + reasoning: list[str] = [] + for cumulative in cumulatives: + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + reasoning_delta, visible_delta = extractor.feed(new_text) + if reasoning_delta: + reasoning.append(reasoning_delta) + if visible_delta: + visible.append(visible_delta) + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + reasoning.append(final_reasoning) + if final_visible: + visible.append(final_visible) + return "".join(visible), "".join(reasoning) + + +def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): + # Parity contract: a reasoning-only reply must reach the client identically + # whether tools are on or off. Both generators stream live then + # resolve to the bare reasoning text; the route's suffix-diff + extractor + # must therefore produce the same (visible, reasoning) split for both. + stream = [ + _sse({"reasoning_content": "The capital"}), + _sse({"reasoning_content": " of France is Paris."}), + _done(), + ] + + tool_backend = _make_backend(monkeypatch, [list(stream)], []) + _patch_monotonic(monkeypatch, [1.0, 2.0, 2.0]) + tool_cumulatives = [ + e["text"] + for e in tool_backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + if e.get("type") == "content" + ] + + no_tool_backend = _make_backend(monkeypatch, [list(stream)], []) + no_tool_cumulatives = [ + y + for y in no_tool_backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + ) + if isinstance(y, str) + ] + + # Both paths stream the reasoning live with the same leading shape. (Raw + # yield lists aren't compared verbatim: the tool path emits a pre-existing + # duplicate trailing event that the route's suffix-diff dedupes.) + assert tool_cumulatives[:3] == no_tool_cumulatives[:3] + # The contract that matters: identical route-level output. + tool_out = _replay_route_reasoning_extractor(tool_cumulatives) + no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) + assert tool_out == no_tool_out + # Pin the shared contract so a change to either path shows up here. + _visible, reasoning = tool_out + assert reasoning == "The capital of France is Paris." + + +def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): + # _drain_silently sibling of the structured-tool close: a bare-JSON tool call + # with a live reasoning prefix must also close before draining, and + # must never leak the drained call text as content. + tool_stream = [ + _sse({"reasoning_content": "Searching now."}), + _sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}), + _done(), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + assert content_before_tool[0] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # The bare-JSON call text was drained, never surfaced as content. + assert not any('"name"' in t for t in content_before_tool) + def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): tool_stream = [ From 01b8085dc2e8fae5d99ee7d236d58b706988773c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 17:10:01 -0700 Subject: [PATCH 18/19] Create ossf.yml (#6952) --- .github/workflows/ossf.yml | 78 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 .github/workflows/ossf.yml diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000000..f9a270540f --- /dev/null +++ b/.github/workflows/ossf.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '21 20 * * 0' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@v3 + with: + sarif_file: results.sarif From 49d1fb38633b2e5b640f7034a3e69734a16396ac Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Wed, 8 Jul 2026 03:08:07 +0200 Subject: [PATCH 19/19] Speed up Studio startup path (#6899) * Speed up Studio startup path * Studio: recheck managed binary executability on preflight cache hit and ignore stale unauthenticated platform fetches Preflight: a matching capability cache fingerprint no longer skips the runnability check when the managed binary's executable bit was cleared (size and mtime unchanged, since chmod bumps ctime not mtime). The cache fast path now confirms the binary is still executable, otherwise it falls back to the CLI help probe so preflight reports Stale and can repair, instead of returning Ready and failing later at backend start. Adds a regression test. Frontend: now that first render is no longer gated on fetchDeviceType, the initial unauthenticated health call can resolve after an authenticated platform fetch. Guard the store so a late unauthenticated or failed non-forced response cannot overwrite an already authoritative device type, tunnel URL, or secure flag. Forced refreshes and the first unauthenticated load are unaffected. * Studio: use access(X_OK) for the preflight cache executability guard A mode bitmask treats any execute bit as launchable, but the executable bits can be set only for another owner or group, or be denied by an ACL, so the current user could still hit PermissionDenied at launch and the cached fast path would wrongly return Ready. access(X_OK) checks real executability for the calling user, so an ownership or permission change correctly falls back to the CLI help probe and the Stale repair path. * Studio: ignore any stale non-forced platform fetch once authoritative Extend the platform store guard so a non-forced health response never overwrites an already authoritative result, not only unauthenticated ones. With a saved token the post-render non-forced request can be authenticated but older than a later forced refresh that already picked up the tunnel URL and secure flag; if that earlier request resolves last it would null those fields. Now any non-forced response is dropped once the store holds a server-reported platform. Forced refreshes and the first authoritative write are unaffected. * Studio: run the managed CLI help probe before trusting the preflight cache Restore running the managed CLI help probe before returning Ready from the desktop capability cache, so a managed install whose venv interpreter or a runtime dependency is broken (while path, size, mtime, and markers are unchanged) is reported Stale for repair rather than proceeding to a backend start that cannot spawn. The capability cache still skips the heavier desktop-capabilities probe on a hit, so a warm cache runs one probe instead of two. Removes the executable-access shortcut, which the help probe now subsumes. --------- Co-authored-by: Daniel Han --- studio/backend/core/inference/orchestrator.py | 8 +- ...t_inference_default_models_non_blocking.py | 42 ++++++ .../frontend/src/components/app-sidebar.tsx | 45 ++++-- studio/frontend/src/config/env.ts | 26 +++- .../frontend/src/features/chat/chat-page.tsx | 43 +++++- .../chat/hooks/use-chat-model-runtime.ts | 12 +- .../src/features/chat/runtime-provider.tsx | 12 +- studio/frontend/src/main.tsx | 16 +-- studio/src-tauri/src/preflight.rs | 132 +++++++++++++++++- studio/src-tauri/src/preflight/managed.rs | 16 +++ 10 files changed, 308 insertions(+), 44 deletions(-) create mode 100644 studio/backend/tests/test_inference_default_models_non_blocking.py diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 19d2230278..cf5d24c367 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -108,13 +108,11 @@ class InferenceOrchestrator: @property def default_models(self) -> list[str]: - # Wait up to 5s for background HF fetch - self._top_models_ready.wait(timeout = 5) top_gguf = self._top_gguf_cache or [] top_hub = self._top_hub_cache or [] - # Curated static defaults first, then HF download-ranked to backfill. - # Send extras so the frontend keeps 4 per category after removing - # downloaded ones. + # Never wait for the remote Hugging Face ranking during startup. Chat's + # first /api/models/list needs curated defaults immediately; the + # background fetch backfills extra choices on later calls. result: list[str] = [] seen: set[str] = set() for m in self._static_models + top_gguf + top_hub: diff --git a/studio/backend/tests/test_inference_default_models_non_blocking.py b/studio/backend/tests/test_inference_default_models_non_blocking.py new file mode 100644 index 0000000000..83a8e7bbfb --- /dev/null +++ b/studio/backend/tests/test_inference_default_models_non_blocking.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Default Chat model metadata must not block on remote Hugging Face discovery.""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch): + sleep_seconds = 2.0 + + def _slow_fetch(self: InferenceOrchestrator) -> None: + time.sleep(sleep_seconds) + self._top_gguf_cache = ["unsloth/slow-GGUF"] + self._top_models_ready.set() + + monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch) + + orchestrator = InferenceOrchestrator() + started = time.monotonic() + defaults = orchestrator.default_models + elapsed = time.monotonic() - started + + assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s" + assert defaults == orchestrator._static_models + assert "unsloth/slow-GGUF" not in defaults + + deadline = time.monotonic() + sleep_seconds + 5 + while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline: + time.sleep(0.05) + + assert "unsloth/slow-GGUF" in orchestrator.default_models diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 2dd0d02515..f59a952b3a 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,11 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { - exportConversationRawJsonl, - exportConversationCsv, - exportConversationShareGPT, -} from "@/features/chat/prompt-storage/prompt-storage-dialog"; import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, @@ -174,6 +169,36 @@ const TestTubeOutlineIcon = TestTube01Icon.slice( 3, ) as typeof TestTube01Icon; + +type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl"; + +const CHAT_EXPORT_OPTIONS: Array<{ + label: string; + format: ConversationExportFormat; +}> = [ + { label: "Raw JSONL", format: "raw-jsonl" }, + { label: "CSV", format: "csv" }, + { label: "ShareGPT JSONL", format: "sharegpt-jsonl" }, +]; + +async function exportConversationByFormat( + threadId: string, + format: ConversationExportFormat, +): Promise { + const exports = await import( + "@/features/chat/prompt-storage/prompt-storage-dialog" + ); + switch (format) { + case "raw-jsonl": + return exports.exportConversationRawJsonl(threadId); + case "csv": + return exports.exportConversationCsv(threadId); + case "sharegpt-jsonl": + return exports.exportConversationShareGPT(threadId); + } +} + + function runStatusDotClass(status: TrainingRunSummary["status"]): string { switch (status) { case "running": @@ -899,11 +924,7 @@ export function AppSidebar() { Export - {[ - { label: "Raw JSONL", fn: exportConversationRawJsonl }, - { label: "CSV", fn: exportConversationCsv }, - { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, - ].map(({ label, fn }) => ( + {CHAT_EXPORT_OPTIONS.map(({ label, format }) => ( { @@ -911,7 +932,9 @@ export function AppSidebar() { const ids = item.type === "single" ? [item.id] : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); - await Promise.all(ids.map((id) => fn(id))); + await Promise.all( + ids.map((id) => exportConversationByFormat(id, format)), + ); } catch { toast.error("Export failed."); } diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index def1b9bad9..63cdd03141 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -54,6 +54,17 @@ export const usePlatformStore = create()((_, get) => ({ isChatOnly: () => get().chatOnly, })); +// Once an authoritative (server-reported) platform has been fetched, a +// non-forced response must not overwrite it. The post-render fetchDeviceType() +// in main.tsx runs before auth is ready and can resolve after the authed +// root-route/provider fetches; such a late write would reset deviceType, +// cloudflareUrl/serverUrl/secure, and fetched, whether it is a browser fallback +// (unauthenticated) or an earlier authenticated request that landed after a +// later forced refresh. Forced refreshes are explicit re-reads, so they still write. +function shouldKeepAuthoritativePlatform(force?: boolean): boolean { + return !force && usePlatformStore.getState().fetched; +} + // `force` re-reads /api/health even if cached, to pick up a late-arriving tunnel URL. export async function fetchDeviceType(options?: { force?: boolean; @@ -81,6 +92,15 @@ export async function fetchDeviceType(options?: { server_url?: string | null; secure?: boolean; }; + // Once the store holds an authoritative (server-reported) platform, a + // non-forced response must not overwrite it. It may be an unauthenticated + // fallback, or an earlier authenticated request that resolved after a + // later forced refresh already picked up device_type and the tunnel + // fields; writing either would reset device type or null the tunnel + // fields. Forced refreshes are explicit re-reads, so they still write. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = data.device_type ?? detectLocalPlatform(); const chatOnly = data.chat_only ?? false; const chatOnlyReason = data.chat_only_reason ?? null; @@ -101,7 +121,11 @@ export async function fetchDeviceType(options?: { } catch { // Backend not ready: use client-side detection so chat-only guard works // on initial load (important for macOS). Keep fetched=false so a later - // call retries against the backend. + // call retries against the backend. But a late non-forced failure must not + // wipe an authoritative platform that already resolved. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = detectLocalPlatform(); const chatOnly = deviceType === "mac"; usePlatformStore.setState({ deviceType, chatOnly, fetched: false }); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index cd7cfc77fc..b155eff780 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -35,7 +35,6 @@ import { useNativeModelDrop, useNativePathLeasesSupported, } from "@/features/native-intents"; -import { ProjectSourcesPanel } from "@/features/rag/components/project-sources-panel"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { isTauri } from "@/lib/api-base"; import { toast } from "@/lib/toast"; @@ -51,7 +50,9 @@ import { Tooltip as TooltipPrimitive } from "radix-ui"; import { type CSSProperties, type ReactElement, + lazy, memo, + Suspense, useCallback, useEffect, useMemo, @@ -134,6 +135,13 @@ import { } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; + +const ProjectSourcesPanel = lazy(() => + import("@/features/rag/components/project-sources-panel").then((module) => ({ + default: module.ProjectSourcesPanel, + })), +); + type LoraCandidate = { id: string; baseModel: string; @@ -1018,7 +1026,15 @@ function ProjectLanding({
{projectTab === "sources" ? ( - + + Loading sources… +
+ } + > + + ) : (
{items.map((item) => { @@ -2246,12 +2262,29 @@ export function ChatPage({ return [...fromLoras, ...localModels]; }, [lorasFromStore, localModels]); - useEffect(() => { - if (getTrainingCompareHandoff()) return; - void refresh(); + const inventoryRefreshStartedRef = useRef(false); + const refreshDeferredModelInventories = useCallback(() => { + inventoryRefreshStartedRef.current = true; + void refresh({ includeLoras: true }); refreshLocalModels(); }, [refresh, refreshLocalModels]); + useEffect(() => { + if (getTrainingCompareHandoff()) return; + void refresh({ includeLoras: false }); + const timeoutId = window.setTimeout(() => { + if (!inventoryRefreshStartedRef.current) { + refreshDeferredModelInventories(); + } + }, 1200); + return () => window.clearTimeout(timeoutId); + }, [refresh, refreshDeferredModelInventories]); + + useEffect(() => { + if (!active || !modelSelectorOpen) return; + refreshDeferredModelInventories(); + }, [active, modelSelectorOpen, refreshDeferredModelInventories]); + useEffect(() => { // ChatPage no longer remounts on navigation, so re-check the handoff whenever // we return to /chat (e.g. from the training progress "compare in chat" action). diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e23d1b0b33..798c1658f9 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -312,14 +312,18 @@ export function useChatModelRuntime() { [], ); - const refresh = useCallback(async (options?: { signal?: AbortSignal }) => { + const refresh = useCallback(async (options?: { + signal?: AbortSignal; + includeLoras?: boolean; + }) => { const signal = options?.signal; + const includeLoras = options?.includeLoras ?? true; setModelsError(null); try { const [listRes, statusRes, lorasRes] = await Promise.all([ listModels(), getInferenceStatus(), - listLoras(), + includeLoras ? listLoras() : Promise.resolve(null), ]); // Cancellation can land while the requests above are in flight. Bail @@ -327,7 +331,9 @@ export function useChatModelRuntime() { if (signal?.aborted) return; setModels(listRes.models.map(toChatModelSummary)); - setLoras(lorasRes.loras.map(toLoraSummary)); + if (lorasRes) { + setLoras(lorasRes.loras.map(toLoraSummary)); + } const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 67980b94c6..b545695f9e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -22,7 +22,6 @@ import { unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime, } from "@assistant-ui/react"; import { createAssistantStream } from "assistant-stream"; -import mammoth from "mammoth"; import { type ReactElement, type ReactNode, @@ -33,7 +32,6 @@ import { useMemo, useRef, } from "react"; -import { extractText, getDocumentProxy } from "unpdf"; import { toast } from "sonner"; import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter"; import { @@ -181,7 +179,10 @@ class PDFAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const buffer = new Uint8Array(await attachment.file.arrayBuffer()); + const [{ extractText, getDocumentProxy }, buffer] = await Promise.all([ + import("unpdf"), + attachment.file.arrayBuffer().then((bytes) => new Uint8Array(bytes)), + ]); const pdf = await getDocumentProxy(buffer); const { text } = await extractText(pdf, { mergePages: true }); return { @@ -298,7 +299,10 @@ class DocxAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const arrayBuffer = await attachment.file.arrayBuffer(); + const [{ default: mammoth }, arrayBuffer] = await Promise.all([ + import("mammoth"), + attachment.file.arrayBuffer(), + ]); const { value } = await mammoth.extractRawText({ arrayBuffer }); return { id: attachment.id, diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx index 0922e764bb..d0ddf2fc6e 100644 --- a/studio/frontend/src/main.tsx +++ b/studio/frontend/src/main.tsx @@ -5,8 +5,8 @@ import { StrictMode } from "react"; import { createRoot } from "react-dom/client"; import "./index.css"; -import { fetchDeviceType } from "./config/env"; import { App } from "./app/app"; +import { fetchDeviceType } from "./config/env"; import { initializeLocale } from "./i18n"; const globalCrypto = globalThis.crypto as Crypto | undefined; @@ -36,10 +36,10 @@ if (!rootElement) { initializeLocale(); -fetchDeviceType().then(() => { - createRoot(rootElement).render( - - - , - ); -}); +createRoot(rootElement).render( + + + , +); + +fetchDeviceType().catch(() => undefined); diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 5a48d26632..7ef5244754 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -498,19 +498,68 @@ mod tests { } #[cfg(unix)] - fn remove_managed_capability_cache() { - let _ = std::fs::remove_file( - dirs::home_dir() + static MANAGED_CAPABILITY_CACHE_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + + #[cfg(unix)] + struct ManagedCapabilityCacheHome { + path: PathBuf, + previous: Option, + } + + #[cfg(unix)] + impl ManagedCapabilityCacheHome { + fn new(test_name: &str) -> Self { + use std::time::{SystemTime, UNIX_EPOCH}; + + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) .unwrap() - .join(".unsloth") - .join("studio") - .join("desktop_capability_cache.json"), - ); + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "unsloth-preflight-cache-{test_name}-{}-{nanos}", + std::process::id() + )); + std::fs::create_dir_all(&path).unwrap(); + let previous = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME"); + std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", &path); + Self { path, previous } + } + } + + #[cfg(unix)] + impl Drop for ManagedCapabilityCacheHome { + fn drop(&mut self) { + if let Some(previous) = &self.previous { + std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", previous); + } else { + std::env::remove_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME"); + } + let _ = std::fs::remove_dir_all(&self.path); + } + } + + #[cfg(unix)] + fn managed_capability_cache_path_for_test() -> PathBuf { + std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") + .map(PathBuf::from) + .or_else(dirs::home_dir) + .unwrap() + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json") + } + + #[cfg(unix)] + fn remove_managed_capability_cache() { + let _ = std::fs::remove_file(managed_capability_cache_path_for_test()); } #[cfg(unix)] #[tokio::test] async fn managed_cli_capability_probe_classifies_core_cases() { + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("core-cases"); remove_managed_capability_cache(); for (name, script, stale_reason) in [ @@ -567,6 +616,75 @@ exit 1 } } + #[cfg(unix)] + #[tokio::test] + async fn managed_cli_capability_help_probe_runs_before_cache() { + use std::fs; + + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("cache-hit"); + + remove_managed_capability_cache(); + // `-h` always succeeds unless `modeh` exists; the desktop-capabilities + // probe always succeeds unless `modecap` exists. Toggling those lets us + // prove the ordering: -h runs on every probe (even a cache hit), while + // the heavier capability probe is skipped once the cache is warm. + let fake = fake_cli( + "cap-cache-hit", + r#"#!/bin/sh +log="$0.calls" +modeh="$0.modeh" +modecap="$0.modecap" +printf '%s\n' "$*" >> "$log" +if [ "$1" = "-h" ]; then + if [ -f "$modeh" ]; then exit 42; fi + exit 0 +fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then + if [ -f "$modecap" ]; then exit 42; fi + printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' + exit 0 +fi +exit 1 +"#, + ); + let bin = fake.bin.clone(); + let calls = bin.with_extension("calls"); + let modeh = bin.with_extension("modeh"); + let modecap = bin.with_extension("modecap"); + + // Cold probe: runs -h and the capability probe, then caches the result. + assert!(matches!( + probe_managed_bin(bin.clone()).await, + ManagedProbe::Ready { .. } + )); + let first_calls = fs::read_to_string(&calls).unwrap(); + assert!(first_calls.contains("-h")); + assert!(first_calls.contains("studio desktop-capabilities --json")); + + // Cache hit: -h still runs, but the capability probe is skipped (breaking + // it via `modecap` proves it is not invoked). + fs::write(&modecap, "broken").unwrap(); + fs::write(&calls, "").unwrap(); + assert!(matches!( + probe_managed_bin(bin.clone()).await, + ManagedProbe::Ready { .. } + )); + assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + + // A non-launchable CLI is caught by the -h probe even with a warm cache: + // preflight reports Stale (for repair) and never trusts the cache. + fs::write(&modeh, "broken").unwrap(); + fs::write(&calls, "").unwrap(); + assert!(matches!( + probe_managed_bin(bin).await, + ManagedProbe::Stale { .. } + )); + assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + + remove_managed_capability_cache(); + } + const EXPECTED_ROOT_ID: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; const OTHER_ROOT_ID: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 57b8365ec5..0d20f271c5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -188,6 +188,16 @@ fn managed_bin_fingerprint(bin: &Path) -> Option { } fn capability_cache_path() -> Option { + #[cfg(test)] + if let Some(home) = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") { + return Some( + PathBuf::from(home) + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json"), + ); + } + dirs::home_dir().map(|home| { home.join(".unsloth") .join("studio") @@ -400,6 +410,12 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool { pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { let started = Instant::now(); + // Always verify the managed CLI actually launches before trusting the cache. + // A matching capability fingerprint does not prove the binary can still run: + // its venv interpreter or a runtime dependency can be broken while the + // path/size/mtime/markers are unchanged, so the -h probe runs first and a + // non-launchable install is reported Stale for repair. The capability cache + // below still skips the heavier desktop-capabilities probe on a hit. if !run_cli_probe(&bin, &["-h"]).await { info!( "Managed preflight: cli unusable for {:?} in {}ms",