From f1fcf0054ca3fbee5fa2d5eb362a0769f670ae1c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 18 May 2026 08:52:09 -0700 Subject: [PATCH 01/32] install scripts: bump unsloth pin to >=2026.5.4 (#5566) PyPI unsloth 2026.5.4 is now live; update install.sh and install.ps1 to require at least that version so fresh installs pull the new release. --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 35951d7ee2..a27af9dd3b 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1285,7 +1285,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.4" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -1293,7 +1293,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.4" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1331,7 +1331,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.4" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -1339,7 +1339,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.4" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1367,7 +1367,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.3" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.4" --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 d59605b6a5..dd4f83fab6 100755 --- a/install.sh +++ b/install.sh @@ -1849,7 +1849,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.3" unsloth-zoo + "unsloth>=2026.5.4" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1857,7 +1857,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.3" unsloth-zoo + "unsloth>=2026.5.4" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2025,7 +2025,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.5.3" unsloth-zoo + "unsloth>=2026.5.4" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -2040,7 +2040,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.5.3" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.4" unsloth-zoo 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..." @@ -2072,7 +2072,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 "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.3" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.4" --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 b7f63d3a9e8db0ec9d460a185ca5269c42988f35 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Tue, 19 May 2026 02:22:26 +0100 Subject: [PATCH 02/32] fix: derive Playwright default model expectation (#5589) --- tests/studio/playwright_chat_ui.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index b081c248d0..73b1a81ae2 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -113,6 +113,23 @@ def fail(m): raise AssertionError(f"[ui] FAIL: {m}") +def expected_default_model(): + override = os.environ.get("EXPECTED_DEFAULT_MODEL") + if override: + return override + + studio_backend = Path(__file__).resolve().parents[2] / "studio" / "backend" + if str(studio_backend) not in sys.path: + sys.path.insert(0, str(studio_backend)) + try: + from core.inference.defaults import DEFAULT_MODELS_GGUF + except Exception as exc: + fail(f"could not import DEFAULT_MODELS_GGUF: {exc}") + if not DEFAULT_MODELS_GGUF: + fail("DEFAULT_MODELS_GGUF is empty") + return DEFAULT_MODELS_GGUF[0] + + def soft_fail(m): """Hard fail in STRICT mode, info-warn otherwise. @@ -475,10 +492,7 @@ with sync_playwright() as p: # list or hides the default would break the first-launch UX, # which is what this assertion guards. step("default_models[0] matches DEFAULT_MODELS_GGUF[0]") - EXPECTED_DEFAULT = os.environ.get( - "EXPECTED_DEFAULT_MODEL", - "unsloth/gemma-4-E2B-it-GGUF", - ) + EXPECTED_DEFAULT = expected_default_model() defaults_resp = evaluate_fetch( page, f"{BASE}/api/models/list", From 27845b1fa28ae6a307d7d562a7806f9824681a33 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 18 May 2026 20:03:55 -0700 Subject: [PATCH 03/32] studio: read Playwright default model from defaults.py without importing it (#5595) * studio: read Playwright default model from defaults.py without importing it The Playwright Chat UI job installs Studio with --no-torch and does not have structlog. Importing core.inference.defaults pulls in core/inference/__init__.py (eager orchestrator -> structlog) and defaults.py's own `import utils.hardware.hardware as hw` (also structlog), so the test died before the first page action. Read DEFAULT_MODELS_GGUF as a literal via ast.literal_eval. Zero side effects, no new test deps, the EXPECTED_DEFAULT_MODEL override still wins. * [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/studio/playwright_chat_ui.py | 42 ++++++++++++++++++++++++------ 1 file changed, 34 insertions(+), 8 deletions(-) diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 73b1a81ae2..72c5f8e673 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -118,16 +118,42 @@ def expected_default_model(): if override: return override - studio_backend = Path(__file__).resolve().parents[2] / "studio" / "backend" - if str(studio_backend) not in sys.path: - sys.path.insert(0, str(studio_backend)) + # Parse DEFAULT_MODELS_GGUF as a literal out of defaults.py instead of + # importing it. The Playwright job installs Studio with --no-torch, so + # the studio.backend.core.inference package init (which eagerly imports + # the orchestrator -> structlog) and defaults.py's own + # `import utils.hardware.hardware as hw` are both unavailable. + import ast + + defaults_path = ( + Path(__file__).resolve().parents[2] + / "studio" + / "backend" + / "core" + / "inference" + / "defaults.py" + ) try: - from core.inference.defaults import DEFAULT_MODELS_GGUF + tree = ast.parse(defaults_path.read_text()) except Exception as exc: - fail(f"could not import DEFAULT_MODELS_GGUF: {exc}") - if not DEFAULT_MODELS_GGUF: - fail("DEFAULT_MODELS_GGUF is empty") - return DEFAULT_MODELS_GGUF[0] + fail(f"could not read {defaults_path}: {exc}") + models = None + for node in tree.body: + if not isinstance(node, ast.Assign): + continue + if not any( + isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF" + for t in node.targets + ): + continue + try: + models = ast.literal_eval(node.value) + except Exception as exc: + fail(f"could not eval DEFAULT_MODELS_GGUF literal: {exc}") + break + if not models: + fail("DEFAULT_MODELS_GGUF not found or empty in defaults.py") + return models[0] def soft_fail(m): From ba710a783a0aeec2961a8dcf673e6c7f9c82f6da Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 May 2026 22:57:43 -0700 Subject: [PATCH 04/32] [pre-commit.ci] pre-commit autoupdate (#5586) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit updates: - [github.com/astral-sh/ruff-pre-commit: v0.15.12 → v0.15.13](https://github.com/astral-sh/ruff-pre-commit/compare/v0.15.12...v0.15.13) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a2a4995d62..9d80fe6ff5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.12 + rev: v0.15.13 hooks: - id: ruff args: From c4908b7929a25fc8f228e7da96164800c4466dbf Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Tue, 19 May 2026 00:55:55 -0700 Subject: [PATCH 05/32] studio: fix toast close-button click and light-mode hover (#5597) Two related issues on the chat toasts: 1. Close X did nothing. The lib/toast.ts wrapper defaulted every toast to `dismissible: false` (originally to keep swipe capture from stealing text selection). In sonner v2, `dismissible: false` makes the close-button onClick a no-op, so the X looked clickable but never dismissed the toast. The Toaster already sets `swipeDirections={[]}` in components/ui/sonner.tsx, so the per-toast swipe workaround is unnecessary and harmful. Replace the wrapper with a thin re-export of sonner. 2. Close X hover collapsed to a near-black circle in light mode. Sonner's default close-button styling uses fixed gray-scale tokens (--gray2 hover, --gray12 text) that ignore the theme attribute. Once the Toaster's inline style overrides --normal-bg with var(--popover), the base background follows the app theme but the hover state does not, so the hover bg lands on a color that has no contrast with the X glyph. Pin both base and hover to theme tokens (--popover, --muted, --popover-foreground, --border) so contrast stays visible in both light and dark modes. Repro: open chat, load any cached model, hover the X on the " loaded" toast in light mode -- before this change the circle turned dark and the click did nothing; after, the circle stays light and the click dismisses the toast. --- studio/frontend/src/index.css | 11 ++++-- studio/frontend/src/lib/toast.ts | 57 ++------------------------------ 2 files changed, 12 insertions(+), 56 deletions(-) diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index dc73112994..7afa87f597 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1172,11 +1172,18 @@ mix-blend-mode: normal; } -/* Override sonner's hardcoded top: 0 on the toast close button. */ +/* Override sonner top: 0 and pin to theme tokens (--gray2 hover ignores data-sonner-theme). */ [data-sonner-toast][data-styled="true"] [data-close-button] { top: 8px !important; + background: var(--popover) !important; + color: var(--popover-foreground) !important; + border-color: var(--border) !important; } -/* Bump the X stroke so it stays visible against dark backgrounds. */ [data-sonner-toast][data-styled="true"] [data-close-button] svg { stroke-width: 2.25; } +[data-sonner-toast][data-styled="true"]:hover [data-close-button]:hover { + background: var(--muted) !important; + color: var(--popover-foreground) !important; + border-color: var(--border) !important; +} diff --git a/studio/frontend/src/lib/toast.ts b/studio/frontend/src/lib/toast.ts index 138ceb4c4d..6b1635b42e 100644 --- a/studio/frontend/src/lib/toast.ts +++ b/studio/frontend/src/lib/toast.ts @@ -1,59 +1,8 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -// sonner `toast` wrapper that defaults `dismissible: false` so swipe -// capture doesn't block text selection. Drop-in for `from "sonner"`. - -import { toast as sonnerToast, type ExternalToast } from "sonner"; - -type AnyFn = (...args: unknown[]) => unknown; - -function withDismissibleFalse(fn: F): F { - return ((...args: unknown[]) => { - // Branch by arity: React-element messages are objects too. - if (args.length <= 1) { - args.push({ dismissible: false } satisfies ExternalToast); - } else { - const lastIdx = args.length - 1; - const last = args[lastIdx]; - if (last && typeof last === "object" && !Array.isArray(last)) { - const opts = last as ExternalToast; - if (!("dismissible" in opts)) { - args[lastIdx] = { dismissible: false, ...opts }; - } - } - } - return fn(...args); - }) as F; -} - -const wrappedCallable = withDismissibleFalse( - sonnerToast as unknown as AnyFn, -) as typeof sonnerToast; - -// `promise(p, data?)` carries `dismissible` at the top of `data`, -// covering loading / success / error states. `dismiss`, `getHistory`, -// `getToasts` take no options. -const wrappedPromise: typeof sonnerToast.promise = ((promise, data) => { - const merged = - data && typeof data === "object" && !("dismissible" in data) - ? { dismissible: false, ...data } - : (data ?? { dismissible: false }); - return sonnerToast.promise(promise, merged); -}) as typeof sonnerToast.promise; - -export const toast: typeof sonnerToast = Object.assign(wrappedCallable, { - success: withDismissibleFalse(sonnerToast.success.bind(sonnerToast) as AnyFn) as typeof sonnerToast.success, - info: withDismissibleFalse(sonnerToast.info.bind(sonnerToast) as AnyFn) as typeof sonnerToast.info, - warning: withDismissibleFalse(sonnerToast.warning.bind(sonnerToast) as AnyFn) as typeof sonnerToast.warning, - error: withDismissibleFalse(sonnerToast.error.bind(sonnerToast) as AnyFn) as typeof sonnerToast.error, - message: withDismissibleFalse(sonnerToast.message.bind(sonnerToast) as AnyFn) as typeof sonnerToast.message, - loading: withDismissibleFalse(sonnerToast.loading.bind(sonnerToast) as AnyFn) as typeof sonnerToast.loading, - custom: withDismissibleFalse(sonnerToast.custom.bind(sonnerToast) as AnyFn) as typeof sonnerToast.custom, - promise: wrappedPromise, - dismiss: sonnerToast.dismiss.bind(sonnerToast) as typeof sonnerToast.dismiss, - getHistory: sonnerToast.getHistory.bind(sonnerToast) as typeof sonnerToast.getHistory, - getToasts: sonnerToast.getToasts.bind(sonnerToast) as typeof sonnerToast.getToasts, -}); +// Re-export of sonner. Swipe blocking lives on the Toaster via +// `swipeDirections={[]}`, so no per-toast dismissible override. +export { toast } from "sonner"; export type { ExternalToast } from "sonner"; From 94026fc8dc8e23423ec6bf232ba9ba6646cba21e Mon Sep 17 00:00:00 2001 From: Junhyuk Lee <58055473+xodn348@users.noreply.github.com> Date: Tue, 19 May 2026 03:05:13 -0500 Subject: [PATCH 06/32] fix(loader): honour HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE in from_pretrained (#5598) Reads HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE in FastLanguageModel.from_pretrained and FastModel.from_pretrained, forcing local_files_only=True so all delegation paths (load_in_4bit, load_in_8bit, full_finetuning, qat_scheme) and direct FastModel callers (FastVisionModel, FastTextModel) honour offline mode. Also gates HF_HUB_ENABLE_HF_TRANSFER in unsloth/dataprep/synthetic.py and adds an early return in get_statistics. Pairs with unslothai/unsloth-zoo#675. Fixes #5316. --- unsloth/dataprep/synthetic.py | 7 ++++++- unsloth/models/_utils.py | 7 +++++++ unsloth/models/loader.py | 19 +++++++++++++++++++ 3 files changed, 32 insertions(+), 1 deletion(-) diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 612c531f47..1be83c3bb9 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -21,7 +21,12 @@ from collections import deque import time import os -os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" +_OFFLINE_VALS = {"1", "true", "yes", "on"} +if not ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_VALS + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_VALS +): + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" import requests import torch import gc diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index a46d1f0c0e..f4313c4518 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1771,6 +1771,13 @@ def get_statistics(local_files_only = False): return if local_files_only: return + # Also skip when HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE are set. + _offline_vals = {"1", "true", "yes", "on"} + if ( + os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline_vals + or os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline_vals + ): + return from huggingface_hub.utils import ( disable_progress_bars, enable_progress_bars, diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index fc91178d88..c10443e289 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -308,6 +308,16 @@ class FastLanguageModel(FastLlamaModel): if is_dist: device_map = distributed_device_map + # Honour offline env vars BEFORE FastModel delegation so 8bit / + # full-finetuning / qat paths also receive local_files_only. + if not kwargs.get("local_files_only", False): + _offline = {"1", "true", "yes", "on"} + if ( + os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline + or os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline + ): + kwargs["local_files_only"] = True + if load_in_8bit or full_finetuning or qat_scheme is not None: return FastModel.from_pretrained( model_name = model_name, @@ -1055,6 +1065,15 @@ class FastModel(FastBaseModel): model_config = None peft_config = None local_files_only = kwargs.get("local_files_only", False) + # Mirror env-var fallback for direct callers (FastVisionModel / FastTextModel). + if not local_files_only: + _offline = {"1", "true", "yes", "on"} + if ( + os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline + or os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline + ): + local_files_only = True + kwargs["local_files_only"] = True try: model_config = AutoConfig.from_pretrained( From 5ce4ab4d54d22af53581b0acf512dfee8d5ce987 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 03:16:05 -0700 Subject: [PATCH 07/32] studio: emit one comma-chained --spec-type for CPU/Mac MTP path (#5575) * studio: emit one comma-chained --spec-type for CPU/Mac MTP path llama-server takes a single --spec-type whose value may be comma-separated to chain implementations (e.g. ngram-mod,draft-mtp). The CPU/Mac MTP branch in LlamaCppBackend.load_model was passing --spec-type twice in the same invocation, which is not the documented chaining mechanism and silently drops one of the two specs depending on llama.cpp's argv handling. Collapse the pair to --spec-type ngram-mod,{mtp_token} and update the stale _extra_args_set_spec_type docstring that claimed llama-server accumulates repeated --spec-type. Update the matching pass-through fixture in test_llama_server_args.py. * studio: align MTP ngram-mod knobs with llama.cpp upstream defaults Two correctness fixes against the llama.cpp server README: 1. The CPU/Mac comma-chained branch was emitting --spec-ngram-mod-n-max 6 with --spec-ngram-mod-n-min 48, which is nonsensical (min > max). Per the upstream default the value is 64. 2. The standalone ngram-mod branch was emitting --spec-ngram-size-n, --draft-min, --draft-max. llama.cpp removed those arg aliases for ngram-mod (they live only on the ngram-simple / map families now); the correct knobs are --spec-ngram-mod-n-match / n-min / n-max. Also refresh the inline comment block to point at the server README rather than the older docs/speculative.md draft- aliases. --- studio/backend/core/inference/llama_cpp.py | 31 +++++++++++-------- .../backend/tests/test_llama_server_args.py | 6 ++-- 2 files changed, 20 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 21f2fe71b5..95a8c26a3a 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -471,8 +471,9 @@ def _is_mtp_model_name( def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: - """User passed --spec-type / --spec-default? llama-server accumulates - repeated --spec-type, so we suppress auto-emit when this is true.""" + """User passed --spec-type / --spec-default? llama-server takes a + single --spec-type (comma-separated to chain), so suppress + auto-emit when this is true.""" if not extra_args: return False for raw in extra_args: @@ -2631,10 +2632,10 @@ class LlamaCppBackend: # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x # - # Params from llama.cpp docs (docs/speculative.md): - # --spec-ngram-size-n 24 (small n not recommended) - # --draft-min 48 --draft-max 64 (MoEs need long drafts; - # dense models can reduce these) + # Params from llama.cpp server README: + # --spec-ngram-mod-n-match 24 (lookup length) + # --spec-ngram-mod-n-min 48 --spec-ngram-mod-n-max 64 + # (MoEs need long drafts; dense models can reduce these) # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md # ref: https://github.com/ggml-org/llama.cpp/pull/19164 # ref: https://github.com/ggml-org/llama.cpp/pull/18471 @@ -2692,20 +2693,22 @@ class LlamaCppBackend: ] ) else: + # CPU/Mac: chain ngram-mod + MTP in one + # comma-separated --spec-type (not repeated). + # ngram-mod knobs match llama.cpp defaults + # (n-match 24, n-min 48, n-max 64). cmd.extend( [ "--spec-type", - mtp_token, + f"ngram-mod,{mtp_token}", "--spec-draft-n-max", "3", - "--spec-type", - "ngram-mod", "--spec-ngram-mod-n-match", "24", "--spec-ngram-mod-n-min", "48", "--spec-ngram-mod-n-max", - "6", + "64", ] ) self._speculative_type = "draft-mtp" @@ -2715,13 +2718,15 @@ class LlamaCppBackend: elif normalized_spec in _valid_spec_types: cmd.extend(["--spec-type", normalized_spec]) if normalized_spec == "ngram-mod": + # llama.cpp defaults; legacy --spec-ngram-size-n + # / --draft-{min,max} were removed for ngram-mod. cmd.extend( [ - "--spec-ngram-size-n", + "--spec-ngram-mod-n-match", "24", - "--draft-min", + "--spec-ngram-mod-n-min", "48", - "--draft-max", + "--spec-ngram-mod-n-max", "64", ] ) diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index f4dabfcf08..68a1c870fb 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -47,17 +47,15 @@ from core.inference.llama_server_args import ( ["--spec-type", "draft-mtp", "--spec-draft-n-max", "6"], [ "--spec-type", - "draft-mtp", + "ngram-mod,draft-mtp", "--spec-draft-n-max", "3", - "--spec-type", - "ngram-mod", "--spec-ngram-mod-n-match", "24", "--spec-ngram-mod-n-min", "48", "--spec-ngram-mod-n-max", - "6", + "64", ], # Reasoning controls ["--reasoning-format", "deepseek"], From 66cfbeac1d6dc26d177244861062b31cee2d65f1 Mon Sep 17 00:00:00 2001 From: swappy <59965507+rycerzes@users.noreply.github.com> Date: Tue, 19 May 2026 16:27:50 +0530 Subject: [PATCH 08/32] Fix loss function not patched for Qwen3.5 models (#5442) * fix: patch loss functions for Qwen3_5ForConditionalGeneration to prevent OOM errors * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Narrow except scope and simplify LOSS_MAPPING sweep Replace bare except Exception with the only two compatibility errors we actually care about so genuine bugs in the sweep surface. Drop the redundant _key != "ForCausalLM" guard since the __name__ predicate already excludes the patched entry (UnslothForCausalLMLoss != ForCausalLMLoss). * [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> Co-authored-by: Daniel Han --- tests/test_import_fixes_drift.py | 67 +++++++++++++++++++++++++++ unsloth/kernels/cross_entropy_loss.py | 15 ++++++ 2 files changed, 82 insertions(+) diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index f90556bf66..099b65d09a 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -545,6 +545,73 @@ def test_transformers_pretrained_model_has_get_input_embeddings(): # =========================================================================== +# =========================================================================== +# transformers LOSS_MAPPING -- patch_loss_functions() coverage +# Regression for https://github.com/unslothai/unsloth/issues/4188: +# Qwen3_5ForConditionalGeneration has loss_type='ForConditionalGeneration', +# a separate LOSS_MAPPING key that was never patched, leaving the model with +# the stock ForCausalLMLoss which does logits.float() and OOMs on <=24 GB GPUs. +# =========================================================================== + + +def _reset_loss_mapping(mapping, saved): + mapping.clear() + mapping.update(saved) + + +def test_patch_loss_functions_covers_conditional_generation(): + """After patch_loss_functions(), every LOSS_MAPPING key that was aliased + to ForCausalLMLoss must also point at the Unsloth kernel -- not just + LOSS_MAPPING['ForCausalLM'].""" + lu = pytest.importorskip("transformers.loss.loss_utils") + cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss") + + saved = dict(lu.LOSS_MAPPING) + try: + cel.patch_loss_functions(torch_compile = False) + + unsloth_loss = lu.LOSS_MAPPING.get("ForCausalLM") + assert unsloth_loss is not None + assert "Unsloth" in str( + unsloth_loss + ), f"LOSS_MAPPING['ForCausalLM'] was not replaced: {unsloth_loss}" + + cg_loss = lu.LOSS_MAPPING.get("ForConditionalGeneration") + assert cg_loss is unsloth_loss, ( + f"LOSS_MAPPING['ForConditionalGeneration'] not patched: {cg_loss}. " + f"Qwen3_5ForConditionalGeneration will silently use the stock " + f"ForCausalLMLoss and OOM at large sequence lengths." + ) + finally: + _reset_loss_mapping(lu.LOSS_MAPPING, saved) + + +def test_patch_loss_functions_does_not_touch_other_loss_types(): + """patch_loss_functions() must not overwrite unrelated loss types + (segmentation, detection, masked-LM, etc.) with the causal-LM kernel.""" + lu = pytest.importorskip("transformers.loss.loss_utils") + cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss") + + non_causal_keys = { + k + for k, v in lu.LOSS_MAPPING.items() + if getattr(v, "__name__", "") != "ForCausalLMLoss" + } + + saved = dict(lu.LOSS_MAPPING) + try: + cel.patch_loss_functions(torch_compile = False) + + unsloth_loss = lu.LOSS_MAPPING.get("ForCausalLM") + for key in non_causal_keys: + assert lu.LOSS_MAPPING.get(key) is not unsloth_loss, ( + f"patch_loss_functions() incorrectly overwrote " + f"LOSS_MAPPING['{key}'] with the Unsloth ForCausalLM kernel." + ) + finally: + _reset_loss_mapping(lu.LOSS_MAPPING, saved) + + def test_accelerate_utils_imports_module_present(): """``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` (import_fixes.py 493-516, 1320-1372). Both reach into accelerate.utils.imports.""" diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index d92229314f..4a8f83ad04 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -461,3 +461,18 @@ if (Version(torch.__version__) < Version("2.4.0")) and not hasattr( # Patch CE Losses in transformers def patch_loss_functions(torch_compile = True): _patch_loss_functions(fast_cross_entropy_loss, torch_compile = torch_compile) + + # Defense-in-depth sweep for LOSS_MAPPING aliases still pointing at the + # stock ForCausalLMLoss (e.g. ForConditionalGeneration for Qwen3.5, + # CsmForConditionalGeneration). unsloth_zoo also does this; remove once + # the floor pin moves past unslothai/unsloth-zoo#656. + try: + import transformers.loss.loss_utils as _lu + + _unsloth_loss = _lu.LOSS_MAPPING.get("ForCausalLM") + if _unsloth_loss is not None: + for _key, _fn in list(_lu.LOSS_MAPPING.items()): + if getattr(_fn, "__name__", "") == "ForCausalLMLoss": + _lu.LOSS_MAPPING[_key] = _unsloth_loss + except (ImportError, AttributeError): + pass From b01a1ba1c2d8840051309290abee8ca97cdd6431 Mon Sep 17 00:00:00 2001 From: alkinun Date: Tue, 19 May 2026 14:36:20 +0300 Subject: [PATCH 09/32] Fix GGUF multi-image chat handling (#5508) Preserves per-turn OpenAI image_url content parts in the standard GGUF /v1/chat/completions path so multi-image chat history keeps each image attached to its original turn. Legacy top-level image_base64 is injected as a synthetic image_url part only when no message-level image exists. Tool use is disabled whenever any GGUF image is present. Fixes #5470. --- studio/backend/routes/inference.py | 88 +++++++------- .../tests/test_openai_tool_passthrough.py | 109 ++++++++++++++++++ 2 files changed, 156 insertions(+), 41 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 607245467c..2ed3315f56 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2254,46 +2254,11 @@ async def openai_chat_completions( detail = "Audio input is not supported for GGUF chat models yet.", ) - # Reject images if this GGUF model doesn't support vision - image_b64 = extracted_image_b64 or payload.image_base64 - if image_b64 and not llama_backend.is_vision: - raise HTTPException( - status_code = 400, - detail = "Image provided but current GGUF model does not support vision.", - ) - - # Convert image to PNG for llama-server (stb_image has limited format support) - if image_b64: - try: - import base64 as _b64 - from io import BytesIO as _BytesIO - from PIL import Image as _Image, UnidentifiedImageError as _UIE - - raw = _b64.b64decode(image_b64) - # Normalize to RGB so PNG encoding succeeds regardless of - # source mode (RGBA, P, L, CMYK, I, F, ...). Previously - # we only converted RGBA, which left CMYK/I/F to raise at - # img.save(PNG). - img = _Image.open(_BytesIO(raw)).convert("RGB") - buf = _BytesIO() - img.save(buf, format = "PNG") - image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii") - except _UIE: - raise HTTPException( - status_code = 400, - detail = "Unsupported or corrupt image format.", - ) - except Exception: - raise HTTPException( - status_code = 400, - detail = "Failed to process image.", - ) - - # Build message list with system prompt prepended - gguf_messages = [] - if system_prompt: - gguf_messages.append({"role": "system", "content": system_prompt}) - gguf_messages.extend(chat_messages) + gguf_messages, has_gguf_image = _openai_messages_for_gguf_chat( + payload, + llama_backend.is_vision, + ) + image_b64 = None cancel_event = threading.Event() @@ -2307,7 +2272,7 @@ async def openai_chat_completions( use_tools = ( _effective_enable_tools(payload) and llama_backend.supports_tools - and not image_b64 + and not has_gguf_image ) if use_tools: @@ -4804,6 +4769,47 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: + """Build llama-server messages for the standard GGUF chat path. + + llama-server accepts OpenAI multimodal content parts directly. Preserve + all per-turn ``image_url`` parts so multi-image chat history keeps each + image attached to its original turn. + """ + messages = _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) + has_message_image = any( + isinstance(msg.get("content"), list) + and any(part.get("type") == "image_url" for part in msg["content"]) + for msg in messages + ) + if payload.image_base64 and not has_message_image: + # Legacy bytes can be any format; the normalizer below sniffs and + # re-encodes to PNG, so the declared mime is rewritten anyway. + image_part = { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{payload.image_base64}", + }, + } + for msg in reversed(messages): + if msg.get("role") != "user": + continue + existing = msg.get("content") + if isinstance(existing, str): + msg["content"] = [{"type": "text", "text": existing}, image_part] + elif isinstance(existing, list): + existing.append(image_part) + else: + msg["content"] = [image_part] + break + else: + messages.append({"role": "user", "content": [image_part]}) + has_image = _normalize_anthropic_openai_images(messages, is_vision) + return messages, has_image + + def _extract_response_format(payload): """Return the ``response_format`` field on an incoming ChatCompletionRequest (or None). The model is declared with ``extra="allow"`` so pydantic stashes diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 638cbc12c8..84f3e41998 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -26,6 +26,7 @@ sys.path.insert(0, _backend) import httpx import pytest +from fastapi import HTTPException from pydantic import ValidationError from models.inference import ( @@ -532,6 +533,7 @@ class TestFriendlyErrorHttpx: from routes.inference import ( # noqa: E402 _drop_empty_assistant_sentinels, + _openai_messages_for_gguf_chat, _openai_messages_for_passthrough, ) @@ -616,3 +618,110 @@ class TestDropEmptyAssistantSentinels: assert roles == ["user", "user"] for m in out: assert m.get("content"), m + + +class TestGgufVisionMessages: + _PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADUlEQVR42mNk" + "+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + + def test_preserves_multiturn_image_parts_on_original_turns(self): + req = ChatCompletionRequest( + model = "default", + image_base64 = self._PNG_B64, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe image one"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{self._PNG_B64}", + }, + }, + ], + }, + {"role": "assistant", "content": "first answer"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "describe image two"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{self._PNG_B64}", + }, + }, + ], + }, + ], + ) + + messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True) + + assert has_image is True + assert messages[0]["content"][0] == { + "type": "text", + "text": "describe image one", + } + assert messages[0]["content"][1]["type"] == "image_url" + assert len(messages[0]["content"]) == 2 + assert messages[2]["content"][0] == { + "type": "text", + "text": "describe image two", + } + assert messages[2]["content"][1]["type"] == "image_url" + assert len(messages[2]["content"]) == 2 + assert isinstance(messages[1]["content"], str) + + # Legacy top-level image_base64 must be ignored when any message-level + # image already exists; otherwise turn 2 ends up with two image parts. + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + image_parts = [p for p in content if p.get("type") == "image_url"] + assert len(image_parts) == 1, msg + + def test_legacy_image_base64_is_injected_when_messages_are_text_only(self): + req = ChatCompletionRequest( + model = "default", + image_base64 = self._PNG_B64, + messages = [{"role": "user", "content": "describe this image"}], + ) + + messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True) + + assert has_image is True + assert messages[0]["content"][0] == { + "type": "text", + "text": "describe this image", + } + assert messages[0]["content"][1]["type"] == "image_url" + assert messages[0]["content"][1]["image_url"]["url"].startswith( + "data:image/png;base64," + ) + + def test_rejects_image_parts_for_text_only_gguf(self): + req = ChatCompletionRequest( + model = "default", + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{self._PNG_B64}", + }, + }, + ], + }, + ], + ) + + with pytest.raises(HTTPException) as exc_info: + _openai_messages_for_gguf_chat(req, is_vision = False) + assert "does not support vision" in str(exc_info.value) From f747108212afe9243edcd6200a07781ac0b2a13f Mon Sep 17 00:00:00 2001 From: alkinun Date: Tue, 19 May 2026 15:06:17 +0300 Subject: [PATCH 10/32] studio: extract tool-call XML parser into a reusable helper module (#5583) Move the inline tool-call XML parser and stripper out of studio/backend/core/inference/llama_cpp.py into a new studio/backend/core/tool_healing.py so external inference servers (llama-server wrappers, llama-swap, custom shims) can reuse the same logic without importing the inference orchestrator, structlog, httpx, or anything from torch / transformers / unsloth. Closes #5502. What this PR does: - New file studio/backend/core/tool_healing.py contains the regex constants (_TOOL_CLOSED_PATS, _TOOL_ALL_PATS, _TC_JSON_START_RE, _TC_FUNC_START_RE, _TC_END_TAG_RE, _TC_FUNC_CLOSE_RE, _TC_PARAM_START_RE, _TC_PARAM_CLOSE_RE), parse_tool_calls_from_text, and strip_tool_call_markup. The regexes and function bodies are byte-for-byte the same as the previous inline implementation in llama_cpp.py; only the @staticmethod decorator and the closure-only `if not auto_heal_tool_calls: return text` short-circuit are dropped (the latter stays in the caller as a fast path when healing is off). - studio/backend/core/inference/llama_cpp.py now imports the regexes and helpers from .tool_healing. LlamaCppBackend._parse_tool_calls_from_text becomes a one-line delegate; the _strip_tool_markup closure keeps the auto_heal_tool_calls fast path and delegates the work. - Helper module imports cleanly without torch, transformers, structlog, httpx, or numpy. studio.backend.core itself is already stdlib-only at import time (lazy __getattr__), so `from studio.backend.core.tool_healing import parse_tool_calls_from_text, strip_tool_call_markup` is the lightweight import path issue #5502 asked for. No behaviour change for existing Studio paths. parse_tool_calls_from_text and strip_tool_call_markup produce the same OpenAI-shape output the old inline code produced for every input. Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 150 ++---------------- studio/backend/core/tool_healing.py | 173 +++++++++++++++++++++ 2 files changed, 187 insertions(+), 136 deletions(-) create mode 100644 studio/backend/core/tool_healing.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 95a8c26a3a..1791f06676 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -28,6 +28,18 @@ from urllib.parse import urlparse import httpx +from core.tool_healing import ( + _TC_END_TAG_RE, + _TC_FUNC_CLOSE_RE, + _TC_FUNC_START_RE, + _TC_JSON_START_RE, + _TC_PARAM_CLOSE_RE, + _TC_PARAM_START_RE, + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + parse_tool_calls_from_text, + strip_tool_call_markup, +) from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, @@ -362,25 +374,6 @@ def _extract_model_size_b(model_id: str): return extract_model_size_b(model_id) -# ── Pre-compiled patterns for tool XML stripping ───────────── -_TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r".*?", re.DOTALL), -] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r".*$", re.DOTALL), - re.compile(r".*$", re.DOTALL), -] - -# ── Pre-compiled patterns for tool-call XML parsing ────────── -_TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_FUNC_START_RE = re.compile(r"\s*") -_TC_END_TAG_RE = re.compile(r"") -_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -_TC_PARAM_START_RE = re.compile(r"\s*") -_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") - - _TOOL_TEMPLATE_MARKERS = ( "{%- if tools %}", "{%- if tools -%}", @@ -3502,119 +3495,7 @@ class LlamaCppBackend: Closing tags (, , ) are all optional since models frequently omit them. """ - tool_calls = [] - - # Pattern 1: JSON inside tags. - # Use balanced-brace extraction that skips braces inside JSON strings. - for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # position of the opening { - depth, i = 0, brace_start - in_string = False - while i < len(content): - ch = content[i] - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 # skip escaped character - continue - if ch == '"': - in_string = False - elif ch == '"': - in_string = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - break - i += 1 - if depth == 0: - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps( - tc["function"]["arguments"] - ) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 2: XML-style value - # All closing tags optional -- models frequently omit , - # , and/or . - if not tool_calls: - # Step 1: Find all positions and extract their bodies. - # Body boundary: use only or next as a boundary because - # code parameter values can contain that literal string. - # After extracting, we trim a trailing if present. - func_starts = list(_TC_FUNC_START_RE.finditer(content)) - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - # Hard boundaries: next - next_func = ( - func_starts[idx + 1].start() - if idx + 1 < len(func_starts) - else len(content) - ) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - # Trim trailing if present (it's the real closing tag) - body = _TC_FUNC_CLOSE_RE.sub("", body) - - # Step 2: Extract parameters from body. - # For single-parameter functions (the common case: code, command, - # query), use body end as the only boundary to avoid false matches - # on inside code strings. - arguments = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - # Single parameter: value is everything from after the tag - # to end of body, trimming any trailing . - pm = param_starts[0] - val = body[pm.end() :] - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - # Value ends at next if present - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - - tc = { - "id": f"call_{len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - - return tool_calls + return parse_tool_calls_from_text(content) @staticmethod def _build_openai_messages( @@ -3993,10 +3874,7 @@ class LlamaCppBackend: def _strip_tool_markup(text: str, *, final: bool = False) -> str: if not auto_heal_tool_calls: return text - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text + return strip_tool_call_markup(text, final = final) # XML prefixes that signal a tool call in content. # Empty when auto_heal is disabled so the buffer never diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py new file mode 100644 index 0000000000..bb61965764 --- /dev/null +++ b/studio/backend/core/tool_healing.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tool-call XML parsing and stripping helpers. + +Extracted verbatim from studio/backend/core/inference/llama_cpp.py so that +external inference servers (llama-server wrappers, llama-swap, custom +shims) can reuse the same logic without importing the inference +orchestrator, structlog, httpx, or the rest of the studio backend. + +The regexes and function bodies are byte-for-byte identical to the +original inline implementation in llama_cpp.py. Any change made here must +preserve that equivalence; tests/python/test_tool_healing_extraction_is_exact.py +verifies it with AST comparison. +""" + +import json +import re + +# Pre-compiled patterns for tool XML stripping. +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + re.compile(r".*?", re.DOTALL), +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r".*$", re.DOTALL), + re.compile(r".*$", re.DOTALL), +] + +# Pre-compiled patterns for tool-call XML parsing. +_TC_JSON_START_RE = re.compile(r"\s*\{") +_TC_FUNC_START_RE = re.compile(r"\s*") +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +_TC_PARAM_START_RE = re.compile(r"\s*") +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + + +def parse_tool_calls_from_text(content: str) -> list[dict]: + """ + Parse tool calls from XML markup in content text. + + Handles formats like: + {"name":"web_search","arguments":{"query":"..."}} + ... + Closing tags (, , ) are all optional + since models frequently omit them. + """ + tool_calls = [] + + # Pattern 1: JSON inside tags. + # Use balanced-brace extraction that skips braces inside JSON strings. + for m in _TC_JSON_START_RE.finditer(content): + brace_start = m.end() - 1 # position of the opening { + depth, i = 0, brace_start + in_string = False + while i < len(content): + ch = content[i] + if in_string: + if ch == "\\" and i + 1 < len(content): + i += 2 # skip escaped character + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + break + i += 1 + if depth == 0: + json_str = content[brace_start : i + 1] + try: + obj = json.loads(json_str) + tc = { + "id": f"call_{len(tool_calls)}", + "type": "function", + "function": { + "name": obj.get("name", ""), + "arguments": obj.get("arguments", {}), + }, + } + if isinstance(tc["function"]["arguments"], dict): + tc["function"]["arguments"] = json.dumps( + tc["function"]["arguments"] + ) + tool_calls.append(tc) + except (json.JSONDecodeError, ValueError): + pass + + # Pattern 2: XML-style value + # All closing tags optional -- models frequently omit , + # , and/or . + if not tool_calls: + # Step 1: Find all positions and extract their bodies. + # Body boundary: use only or next as a boundary because + # code parameter values can contain that literal string. + # After extracting, we trim a trailing if present. + func_starts = list(_TC_FUNC_START_RE.finditer(content)) + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + # Hard boundaries: next + next_func = ( + func_starts[idx + 1].start() + if idx + 1 < len(func_starts) + else len(content) + ) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = content[body_start:body_end] + # Trim trailing if present (it's the real closing tag) + body = _TC_FUNC_CLOSE_RE.sub("", body) + + # Step 2: Extract parameters from body. + # For single-parameter functions (the common case: code, command, + # query), use body end as the only boundary to avoid false matches + # on inside code strings. + arguments = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + # Single parameter: value is everything from after the tag + # to end of body, trimming any trailing . + pm = param_starts[0] + val = body[pm.end() :] + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = val.strip() + else: + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + # Value ends at next if present + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = val.strip() + + tc = { + "id": f"call_{len(tool_calls)}", + "type": "function", + "function": { + "name": func_name, + "arguments": json.dumps(arguments), + }, + } + tool_calls.append(tc) + return tool_calls + + +def strip_tool_call_markup(text: str, *, final: bool = False) -> str: + """Strip tool-call XML markup from text. + + When ``final`` is False, only fully closed tool-call blocks are removed. + When ``final`` is True, trailing incomplete tool-call blocks are removed + too, and the result is stripped of surrounding whitespace. + """ + patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + for pat in patterns: + text = pat.sub("", text) + return text.strip() if final else text From e4edd34e3c0267d682ae23bb21f51d5839108d0e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:48:46 -0700 Subject: [PATCH 11/32] studio/frontend: reconcile stale must_change_password localStorage flag (#5576) * studio/frontend: reconcile stale must_change_password localStorage flag The client OR's a localStorage flag against /api/auth/status everywhere it gates change-password routing, but never clears the flag when the server flips requires_password_change back to false. A user whose default admin password was already rotated (change-password from another browser, the CLI reset-password command, or a recreated auth DB) keeps that flag, so: 1. requirePasswordChangeFlow lets them sit on the change-password route. 2. Back to login bounces via requireGuest, hasActiveSession (which only checks key presence, not validity), then getPostAuthRoute, which sends the user back because the flag is still set. End result: the user is pinned on change-password and cannot escape without clearing localStorage by hand. Fix the three places that compare server status to the flag: - auth-guards.ts fetchAuthStatus: clear the local flag whenever the server reports requires_password_change = false. - auth-guards.ts requireGuest: call fetchAuthStatus before routing so a stale flag cannot decide getPostAuthRoute. - auth-form.tsx initializeAuthForm: same reconcile inside the page so the change-password page redirects to login as soon as it loads when the server no longer requires a change. - api.ts redirectToAuth: same reconcile in the fetch wrapper's auth redirect. After the reconcile the redundant mustChangePassword() OR clauses are no longer load bearing for the change-password gates; the server's fetchAuthStatus is now the single source of truth. * shorter comments around auth-status reconcile * studio/auth: make localStorage reconcile bidirectional --- studio/frontend/src/app/auth-guards.ts | 10 +++++++++- studio/frontend/src/features/auth/api.ts | 6 +++++- .../src/features/auth/components/auth-form.tsx | 11 ++++++++--- 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/studio/frontend/src/app/auth-guards.ts b/studio/frontend/src/app/auth-guards.ts index 52230f0b6f..6849f380b8 100644 --- a/studio/frontend/src/app/auth-guards.ts +++ b/studio/frontend/src/app/auth-guards.ts @@ -9,6 +9,7 @@ import { hasRefreshToken, mustChangePassword, refreshSession, + setMustChangePassword, } from "@/features/auth"; async function hasActiveSession(): Promise { @@ -26,7 +27,12 @@ async function fetchAuthStatus(): Promise { try { const res = await fetch(apiUrl("/api/auth/status")); if (!res.ok) return { initialized: true, requires_password_change: mustChangePassword() }; - return (await res.json()) as AuthStatus; + const status = (await res.json()) as AuthStatus; + // Server truth wins; keep localStorage in sync both ways. + if (status.requires_password_change !== mustChangePassword()) { + setMustChangePassword(status.requires_password_change); + } + return status; } catch { return { initialized: true, requires_password_change: mustChangePassword() }; } @@ -61,6 +67,8 @@ export async function requireGuest(): Promise { throw redirect({ to: "/chat" }); } if (!(await hasActiveSession())) return; + // Reconcile localStorage before routing. + await fetchAuthStatus(); throw redirect({ to: getPostAuthRoute() }); } diff --git a/studio/frontend/src/features/auth/api.ts b/studio/frontend/src/features/auth/api.ts index 98c5757d2a..cc5749a923 100644 --- a/studio/frontend/src/features/auth/api.ts +++ b/studio/frontend/src/features/auth/api.ts @@ -65,7 +65,11 @@ async function redirectToAuth(): Promise { const res = await fetch(apiUrl("/api/auth/status")); if (res.ok) { const data = (await res.json()) as { requires_password_change: boolean }; - if (data.requires_password_change || mustChangePassword()) target = "/change-password"; + // Server truth wins; keep localStorage in sync both ways. + if (data.requires_password_change !== mustChangePassword()) { + setMustChangePassword(data.requires_password_change); + } + if (data.requires_password_change) target = "/change-password"; } } catch { // Fall through to /login on error diff --git a/studio/frontend/src/features/auth/components/auth-form.tsx b/studio/frontend/src/features/auth/components/auth-form.tsx index a10c77e9fa..3ac97929d8 100644 --- a/studio/frontend/src/features/auth/components/auth-form.tsx +++ b/studio/frontend/src/features/auth/components/auth-form.tsx @@ -105,12 +105,17 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null { setInitialized(result.initialized); setRequiresPasswordChange(result.requires_password_change); + // Server truth wins; keep localStorage in sync both ways. + if (result.requires_password_change !== mustChangePassword()) { + setMustChangePassword(result.requires_password_change); + } + // Redirect between login ↔ change-password based on server state if (mode === "login" && result.requires_password_change) { navigate({ to: "/change-password" }); return; } - if (mode === "change-password" && !result.requires_password_change && !mustChangePassword()) { + if (mode === "change-password" && !result.requires_password_change) { navigate({ to: "/login" }); return; } @@ -163,14 +168,14 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null { const blockedByState = initialized === false || (mode === "login" && requiresPasswordChange) || - (mode === "change-password" && !requiresPasswordChange && !mustChangePassword()); + (mode === "change-password" && !requiresPasswordChange); let helperText: string | null = null; if (initialized === false) { helperText = "Auth is still bootstrapping the default admin account."; } else if (isLoginMode && requiresPasswordChange) { helperText = "Sign in once with the seeded credentials to change the password."; - } else if (!isLoginMode && !requiresPasswordChange && !mustChangePassword()) { + } else if (!isLoginMode && !requiresPasswordChange) { helperText = "Password already updated. Use the login screen."; } const title = isLoginMode ? "Welcome back" : "Setup your account"; From f7e8a85d32df76681a886a823e66eb303391d7ab Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:48:59 -0700 Subject: [PATCH 12/32] studio/frontend: cap auto-load cascade attempts (#5578) * studio/frontend: cap auto-load cascade attempts autoLoadSmallestModel walks every cached GGUF and safetensors repo with a try/catch + continue, so a folder of broken caches (missing files, stale llama.cpp prebuilt, GPU OOM) can fire dozens of failing POST /api/inference/load calls in a row. Each call costs ~5 seconds (HF metadata probe + DNS guard inside inference.py), so the user sees a runaway sequence of request_completed log lines after sending one message that needed an auto-load. Cap the total loadModel calls inside autoLoadSmallestModel at 3 (GGUF cascade plus safetensors fallback share the same counter). Caching that fails three times in a row is almost certainly an environment problem, not "we haven't found the working one yet"; the default-Gemma download path still runs. No behavior change on the happy path: success returns after the first hit exactly like today, and the trust-remote-code skip path does not consume an attempt slot. * shorter comment on auto-load cap * studio chat: extend autoload cap to default Gemma fallback Cached cascade respected MAX_AUTO_LOAD_ATTEMPTS but the default-Gemma download path skipped the budget, so a broken cache could still emit a fourth /api/inference/load. Gate the fallback on the same cap (and bump loadAttempts when we do call loadModel) so the total cross-path budget is 3, matching the cap's intent. --- .../src/features/chat/api/chat-adapter.ts | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 61d71b641a..f87229d78f 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -437,6 +437,9 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { * without selecting one. Prefers GGUF (picks smallest cached variant), * falls back to smallest cached safetensors model. */ +// Cap cascade so broken cached repos can't spam /api/inference/load. +const MAX_AUTO_LOAD_ATTEMPTS = 3; + async function autoLoadSmallestModel(): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean; @@ -451,6 +454,7 @@ async function autoLoadSmallestModel(): Promise<{ }); let blockedByTrustRemoteCode = false; let hadNonTrustFailure = false; + let loadAttempts = 0; async function canAutoLoad(payload: { model_path: string; @@ -481,6 +485,7 @@ async function autoLoadSmallestModel(): Promise<{ if (ggufRepos.length > 0) { const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes); for (const repo of sorted) { + if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { const variants = await listGgufVariants(repo.repo_id); const downloaded = variants.variants @@ -498,6 +503,7 @@ async function autoLoadSmallestModel(): Promise<{ ) { continue; } + loadAttempts += 1; const loadResp = await loadModel({ model_path: repo.repo_id, hf_token: hfToken, @@ -560,6 +566,7 @@ async function autoLoadSmallestModel(): Promise<{ if (modelRepos.length > 0) { const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes); for (const repo of sorted) { + if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { if ( !(await canAutoLoad({ @@ -571,6 +578,7 @@ async function autoLoadSmallestModel(): Promise<{ ) { continue; } + loadAttempts += 1; const sfLoadResp = await loadModel({ model_path: repo.repo_id, hf_token: hfToken, @@ -616,6 +624,17 @@ async function autoLoadSmallestModel(): Promise<{ } } + // Cap also gates the default download so the total /api/inference/load + // budget across cached + fallback is MAX_AUTO_LOAD_ATTEMPTS, not +1. + if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) { + toast.dismiss(toastId); + return { + loaded: false, + blockedByTrustRemoteCode: + blockedByTrustRemoteCode && !hadNonTrustFailure, + }; + } + // No cached models found — try downloading a small default GGUF toast("Downloading a small model…", { id: toastId, @@ -634,6 +653,7 @@ async function autoLoadSmallestModel(): Promise<{ toast.dismiss(toastId); return { loaded: false, blockedByTrustRemoteCode }; } + loadAttempts += 1; const loadResp = await loadModel({ model_path: "unsloth/gemma-4-E2B-it-GGUF", hf_token: hfToken, From d1681ea158c7de921e36e4208bce2664b33caa11 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:49:10 -0700 Subject: [PATCH 13/32] studio: regenerate desktop launcher on `unsloth studio update` (macOS + Linux + Windows) (#5577) * studio: regenerate desktop launcher on `unsloth studio update` Today `unsloth studio update` only mutates the venv. The macOS .app bundle, the Linux .desktop file, and the shared launch-studio.sh stub bake their paths and `studio_install_id` at install time and never refresh. Users who update an existing Studio install report the Dock / Applications icon still pointing at the old launcher; only a fresh `curl ... install.sh | sh` fixes it because that path re-enters install.sh's create_studio_shortcuts. Wire the same logic into the update path: - install.sh: add --shortcuts-only. Skips the heavy install steps, resolves STUDIO_HOME / OS / DATA_DIR through the existing _resolve_studio_destinations + platform detection, then calls create_studio_shortcuts and exits. - unsloth_cli/commands/studio.py: after setup.sh succeeds, call install.sh with --shortcuts-only. Prefers a local checkout's install.sh (when STUDIO_LOCAL_REPO is set) or one shipped under _PACKAGE_ROOT, and falls back to fetching the upstream installer from https://unsloth.ai/install.sh for PyPI-installed users (the wheel does not ship install.sh). Net effect: `unsloth studio update` now refreshes the macOS .app stub, launcher script, studio.conf, and Linux .desktop entry on every update, so the desktop icon stays in sync with the venv that setup.sh just updated. Env-override and Tauri modes keep their existing behavior (no persistent menu shortcuts, but the launch-studio.sh is still regenerated). Windows is unchanged here; setup.ps1 already handles its own Start Menu / Desktop .lnk creation on update. * studio: also regenerate Windows .lnk shortcuts on update Mirror the macOS fix: install.ps1 gains --shortcuts-only that short-circuits to New-StudioShortcuts, and unsloth studio update calls it after setup.ps1 the same way it now does on macOS / Linux. PyPI installs do not ship install.ps1, so the Python helper fetches the upstream script from https://unsloth.ai/install.ps1 and pipes it into powershell.exe -Command - with an explicit Install-UnslothStudio call appended (irm | iex relies on the trailing @args, which is empty when launched from stdin). setup.ps1 alone never recreates the Start Menu / Desktop .lnk targets or the launch-studio.{ps1,vbs} scripts, so without this update users on Windows hit the same stale-icon regression that triggered the macOS PR. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: rename unsloth.exe to .deleteme before update on Windows Pip's editable reinstall calls uninstall first, which deletes every RECORD entry. unsloth.exe is one of them, and Windows refuses to delete a file whose image is mapped into the running process tree. The first unsloth studio update after install therefore fails with: OSError: [WinError 32] The process cannot access the file because it is being used by another process: ...\Scripts\unsloth.exe Windows does allow renaming an in-use exe, so move it aside before _run_setup_script kicks pip. pip then drops a fresh unsloth.exe at the original path; the *.exe.deleteme left behind is cleaned up at the start of the next update once the previous shim has exited. * studio: rename unsloth.exe from setup.ps1 to reliably bypass exe lock * studio: print python -m workaround when Windows exe lock blocks update * studio: use python -c hint (unsloth_cli has no __main__) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install.sh: reshape --shortcuts-only Tauri guard to pass exit-order test * shorter comments in update / launcher regen logic * studio update: env-mode passthrough + non-silent shortcuts-only error * studio update: address codex/gemini PR review - Strip install.ps1's `Install-UnslothStudio @args` auto-invoke before appending an explicit `--shortcuts-only` call so PyPI Windows installs don't re-run the full installer over stdin. - subprocess.run(input=wrapper, ...) now uses encoding="utf-8" so box drawing chars in install.ps1 don't UnicodeEncodeError on CP1252. - Wrap _run_setup_script in try/except to restore unsloth.exe from .deleteme if setup fails, and mirror that rollback inside setup.ps1 when install_python_stack.py exits non-zero. - Capture subprocess return codes in _refresh_desktop_shortcuts and echo a one-line warning on non-zero so silent stale-shortcut failures surface. - Drop --local from the Windows lock-recovery hint so users on PyPI installs don't accidentally switch into editable-checkout mode. - Quote $VENV_ABS_BIN/unsloth in the install.sh shortcuts-only error so paths with spaces print legibly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio update: harden Windows refresh per multi-reviewer pass - PowerShell stdin path now writes the wrapper to a UTF-8 BOM tempfile and runs it via `-File`. `powershell.exe -Command -` decodes stdin with the OEM code page, which mangles box-drawing chars in the fetched install.ps1; -File reads the BOM and decodes UTF-8 cleanly. - _restore_self_exe_lock_windows now treats a zero-byte unsloth.exe as a partial-write and prefers the .deleteme copy. setup.ps1 mirrors the same check. - _release_self_exe_lock_windows uses os.replace for atomic overwrite so a stale .deleteme from an aborted prior update doesn't break the rename. - Lock-recovery hint mentions that --local should be re-added when the user installed from a repo checkout. * studio update: respect Tauri context and tidy Windows .deleteme Tauri's update.rs spawns `unsloth studio update`; without a signal, the CLI's _refresh_desktop_shortcuts would call install.{sh,ps1} --shortcuts-only and create duplicate ~/Applications/Unsloth Studio.app (or .desktop / .lnk) entries that collide with the Tauri bundle. - update.rs now sets UNSLOTH_TAURI_UPDATE=1 on the spawned child. - studio.py's update() skips _refresh_desktop_shortcuts when that env var is set; Tauri owns its own bundle entries. - After a successful Windows update, drop the .deleteme orphan so repeated updates don't accumulate stale binaries that could later be promoted by _restore_self_exe_lock_windows on a cross-version failure. - Tempfile for the PyPI-fallback PowerShell path now uses an unsloth-studio-refresh- prefix so AV/EDR rules and user greps can identify it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio update: drop obsolete WinError 32 hint, echo Tauri skip The rename trick in _release_self_exe_lock_windows + setup.ps1's restore now handle the .exe-lock case in-flow; the printed hint suggested re-running update via venv python, but that just re-enters the same update() and hits the same failure if the rename didn't help. Removing the misleading hint and its helper. Also surface a one-line typer.echo when refresh is skipped under UNSLOTH_TAURI_UPDATE so --verbose logs make the branch visible. * [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> --- install.ps1 | 15 ++ install.sh | 16 +++ studio/setup.ps1 | 37 +++++ studio/src-tauri/src/update.rs | 4 + unsloth_cli/commands/studio.py | 253 ++++++++++++++++++++++++++++++++- 5 files changed, 324 insertions(+), 1 deletion(-) diff --git a/install.ps1 b/install.ps1 index a27af9dd3b..5e3d4b6a50 100644 --- a/install.ps1 +++ b/install.ps1 @@ -92,6 +92,7 @@ function Install-UnslothStudio { $RepoRoot = "" $TauriMode = $false $SkipTorch = $false + $ShortcutsOnly = $false $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -100,6 +101,7 @@ function Install-UnslothStudio { "--no-torch" { $SkipTorch = $true } "--verbose" { $script:UnslothVerbose = $true } "-v" { $script:UnslothVerbose = $true } + "--shortcuts-only" { $ShortcutsOnly = $true } "--package" { $i++ if ($i -ge $argList.Count) { @@ -871,6 +873,19 @@ shell.Run cmd, 0, False } } + # Regen .lnk + launcher only; used by `unsloth studio update`. + if ($ShortcutsOnly) { + if ($TauriMode) { return } + $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" + if (-not (Test-Path -LiteralPath $UnslothExe)) { + Write-Host "[ERROR] unsloth.exe missing at $UnslothExe; run install.ps1 first." -ForegroundColor Red + # throw (not Exit-InstallFailure) so non-Tauri callers see rc != 0. + throw "unsloth.exe missing" + } + New-StudioShortcuts -UnslothExePath $UnslothExe + return + } + # ── Check winget ── Write-TauriLog "STEP" "Checking system dependencies" if (-not (Get-Command winget -ErrorAction SilentlyContinue)) { diff --git a/install.sh b/install.sh index dd4f83fab6..cfd76fa945 100755 --- a/install.sh +++ b/install.sh @@ -45,6 +45,7 @@ TAURI_MODE=false _USER_PYTHON="" _NO_TORCH_FLAG=false _VERBOSE=false +_SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false for arg in "$@"; do @@ -65,6 +66,7 @@ for arg in "$@"; do --python) _next_is_python=true ;; --no-torch) _NO_TORCH_FLAG=true ;; --verbose|-v) _VERBOSE=true ;; + --shortcuts-only) _SHORTCUTS_ONLY=true ;; esac done @@ -1233,6 +1235,20 @@ elif grep -qi microsoft /proc/version 2>/dev/null; then fi step "platform" "$OS" +# Regen launcher/shortcuts only; used by `unsloth studio update`. +if [ "$_SHORTCUTS_ONLY" = true ]; then + # Tauri owns its own shortcuts. + if [ "$TAURI_MODE" != true ]; then + VENV_ABS_BIN="$VENV_DIR/bin" + if [ ! -x "$VENV_ABS_BIN/unsloth" ]; then + echo "ERROR: unsloth binary missing at '$VENV_ABS_BIN/unsloth'; run install.sh first." >&2 + exit 1 + fi + create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" + fi + exit 0 +fi + # ── Architecture detection & Python version ── _ARCH=$(uname -m) MAC_INTEL=false diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 40788a0ecb..16df87bbd5 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1842,6 +1842,21 @@ if ($CuTag -eq "cpu") { } } +# Rename running unsloth.exe so pip can replace it (Windows refuses to delete a mapped .exe). +$VenvScriptsDir = Join-Path $VenvDir "Scripts" +$RunningUnslothExe = Join-Path $VenvScriptsDir "unsloth.exe" +if (Test-Path -LiteralPath $RunningUnslothExe -PathType Leaf) { + $StaleUnslothExe = "$RunningUnslothExe.deleteme" + if (Test-Path -LiteralPath $StaleUnslothExe) { + Remove-Item -LiteralPath $StaleUnslothExe -Force -ErrorAction SilentlyContinue + } + try { + Rename-Item -LiteralPath $RunningUnslothExe -NewName "unsloth.exe.deleteme" -Force -ErrorAction Stop + } catch { + substep "could not rename unsloth.exe ($($_.Exception.Message)); pip may fail with WinError 32" "Yellow" + } +} + # Ordered heavy dependency installation -- shared cross-platform script substep "running ordered dependency installation..." python "$PSScriptRoot\install_python_stack.py" @@ -1851,6 +1866,28 @@ $ErrorActionPreference = $prevEAP if ($stackExit -ne 0) { Write-Host "[FAILED] Python dependency installation failed (exit code $stackExit)" -ForegroundColor Red Write-Host " Re-run the installer or check the error above for details." -ForegroundColor Red + # Restore the pre-rename unsloth.exe so the user keeps a working CLI. + # Treat a zero-byte exe as "pip half-wrote a broken binary" -- prefer the + # stale-but-working copy in .deleteme. + if (Test-Path -LiteralPath "$RunningUnslothExe.deleteme") { + $needRestore = -not (Test-Path -LiteralPath $RunningUnslothExe) + if (-not $needRestore) { + try { + $needRestore = (Get-Item -LiteralPath $RunningUnslothExe -ErrorAction Stop).Length -eq 0 + } catch { $needRestore = $true } + } + if ($needRestore) { + try { + if (Test-Path -LiteralPath $RunningUnslothExe) { + Remove-Item -LiteralPath $RunningUnslothExe -Force -ErrorAction SilentlyContinue + } + Rename-Item -LiteralPath "$RunningUnslothExe.deleteme" -NewName "unsloth.exe" -Force -ErrorAction Stop + substep "restored unsloth.exe after failed install" + } catch { + substep "could not restore unsloth.exe ($($_.Exception.Message))" "Yellow" + } + } + } exit 1 } diff --git a/studio/src-tauri/src/update.rs b/studio/src-tauri/src/update.rs index 390ea5878a..6136c5ea81 100644 --- a/studio/src-tauri/src/update.rs +++ b/studio/src-tauri/src/update.rs @@ -65,6 +65,10 @@ fn spawn_update( // the same install the desktop app uses, not an inherited custom root. cmd.env_remove("UNSLOTH_STUDIO_HOME"); cmd.env_remove("STUDIO_HOME"); + // Signal to unsloth_cli that this update was initiated by the Tauri + // desktop bundle so it skips re-creating CLI launchers/.app/.desktop + // shortcuts (Tauri owns its own bundle entries). + cmd.env("UNSLOTH_TAURI_UPDATE", "1"); #[cfg(windows)] let mut child: Box = { diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 714ace4533..67395a8378 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -6,6 +6,7 @@ import hashlib import json import os import platform +import re import secrets import sqlite3 import subprocess @@ -13,6 +14,8 @@ import sys import tempfile import time import types +import urllib.error +import urllib.request from datetime import datetime, timezone from pathlib import Path from typing import List, Optional @@ -1064,6 +1067,170 @@ def _run_setup_script(*, verbose: bool = False) -> None: raise typer.Exit(result.returncode) +_INSTALLER_URL_BASH = "https://unsloth.ai/install.sh" +_INSTALLER_URL_PWSH = "https://unsloth.ai/install.ps1" + + +def _refresh_desktop_shortcuts(*, verbose: bool = False) -> None: + """Re-run installer with --shortcuts-only to refresh launchers post-update.""" + env = {**os.environ} + if verbose: + env["UNSLOTH_VERBOSE"] = "1" + + is_windows = platform.system() == "Windows" + installer_name = "install.ps1" if is_windows else "install.sh" + installer_url = _INSTALLER_URL_PWSH if is_windows else _INSTALLER_URL_BASH + + # Prefer local checkout, fall back to package dir, then network fetch. + local_repo = (os.environ.get("STUDIO_LOCAL_REPO") or "").strip() + candidates: list[Path] = [] + if local_repo: + candidates.append(Path(local_repo) / installer_name) + candidates.append(_PACKAGE_ROOT / installer_name) + + args = ["--shortcuts-only"] + if verbose: + args.append("--verbose") + + if is_windows: + ps_argv: list[str] = ["powershell.exe"] + if _should_hide_windows_subprocesses(): + ps_argv.extend( + ["-NoLogo", "-NoProfile", "-NonInteractive", "-WindowStyle", "Hidden"] + ) + + for script in candidates: + try: + if script.is_file(): + quoted = str(script).replace("'", "''") + argv = list(ps_argv) + argv.extend( + [ + "-ExecutionPolicy", + "Bypass", + "-Command", + f"& '{quoted}' {' '.join(args)} *>&1", + ] + ) + result = subprocess.run( + argv, + env = env, + check = False, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher install.ps1 exited {result.returncode}" + ) + return + except OSError: + continue + + # PyPI installs lack install.ps1: fetch + pipe to powershell stdin. + try: + request = urllib.request.Request( + installer_url, headers = {"User-Agent": "unsloth-studio-update"} + ) + with urllib.request.urlopen(request, timeout = 30) as response: + installer = response.read().decode("utf-8", errors = "replace") + except (urllib.error.URLError, TimeoutError, OSError) as exc: + typer.echo( + f" refresh-launcher skipped: could not fetch {installer_url} ({exc})" + ) + return + + # install.ps1 auto-invokes `Install-UnslothStudio @args` at EOF; over + # stdin `$args` is empty so that triggers the full installer flow + # (deps, venv, prompts) before our shortcuts-only call. Strip it. + installer = re.sub( + r"(?m)^[ \t]*Install-UnslothStudio[ \t]+@args[ \t]*\r?\n?", + "", + installer, + ) + # stdin-piped scripts have empty $args, so call Install-UnslothStudio explicitly. + marker_args = " ".join(args) + wrapper = installer + f"\nInstall-UnslothStudio {marker_args}\n" + + # Write to a UTF-8 BOM tempfile and use -File rather than -Command -. + # `powershell.exe -Command -` reads stdin via [Console]::InputEncoding + # (CP1252/OEM on most Windows boxes), which mangles box-drawing chars + # in install.ps1. -File reads the BOM and decodes correctly. The + # prefix gives AV/EDR engines (and grep'ing users) a clear identity. + ps1_fd, ps1_path = tempfile.mkstemp( + prefix = "unsloth-studio-refresh-", + suffix = ".ps1", + ) + try: + with os.fdopen(ps1_fd, "wb") as fh: + fh.write(b"\xef\xbb\xbf" + wrapper.encode("utf-8")) + argv = list(ps_argv) + argv.extend(["-ExecutionPolicy", "Bypass", "-File", ps1_path]) + try: + result = subprocess.run( + argv, + env = env, + check = False, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher fetched install.ps1 exited {result.returncode}" + ) + except OSError as exc: + typer.echo( + f" refresh-launcher skipped: powershell exec failed ({exc})" + ) + finally: + try: + os.unlink(ps1_path) + except OSError: + pass + return + + for script in candidates: + try: + if script.is_file(): + result = subprocess.run( + ["bash", str(script), *args], + env = env, + check = False, + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher install.sh exited {result.returncode}" + ) + return + except OSError: + continue + + # PyPI installs lack install.sh: fetch upstream. + try: + request = urllib.request.Request( + installer_url, headers = {"User-Agent": "unsloth-studio-update"} + ) + with urllib.request.urlopen(request, timeout = 30) as response: + installer = response.read() + except (urllib.error.URLError, TimeoutError, OSError) as exc: + typer.echo( + f" refresh-launcher skipped: could not fetch {installer_url} ({exc})" + ) + return + + try: + result = subprocess.run( + ["bash", "-s", "--", *args], + input = installer, + env = env, + check = False, + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher fetched install.sh exited {result.returncode}" + ) + except OSError as exc: + typer.echo(f" refresh-launcher skipped: bash exec failed ({exc})") + + @studio_app.command(hidden = True) def setup( verbose: bool = typer.Option( @@ -1093,6 +1260,9 @@ def update( ), ): """Update Unsloth Studio dependencies and rebuild.""" + # Re-export UNSLOTH_STUDIO_HOME for env-mode installs so the refresh + # subprocess resolves the same install root the user originally chose. + _ensure_studio_env_exported() # Ensure SKIP_STUDIO_BASE is not inherited from a parent install.ps1 session os.environ.pop("SKIP_STUDIO_BASE", None) os.environ["STUDIO_PACKAGE_NAME"] = package @@ -1105,7 +1275,88 @@ def update( else: os.environ["STUDIO_LOCAL_INSTALL"] = "0" os.environ.pop("STUDIO_LOCAL_REPO", None) - _run_setup_script(verbose = verbose) + _release_self_exe_lock_windows() + try: + _run_setup_script(verbose = verbose) + except BaseException: + # Restore unsloth.exe from .deleteme if setup failed before pip + # produced a replacement; otherwise the user has no CLI for recovery. + _restore_self_exe_lock_windows() + raise + # On Windows clear the .deleteme orphan now that pip wrote a fresh + # unsloth.exe; on next update os.replace would overwrite it anyway, + # but leaving a stale binary around invites cross-version restore + # confusion from _restore_self_exe_lock_windows. + _cleanup_self_exe_lock_windows() + # Tauri desktop owns its own bundle entries; skip CLI launcher refresh + # so a Tauri-initiated update doesn't create duplicate shortcuts. + if os.environ.get("UNSLOTH_TAURI_UPDATE") == "1": + if verbose: + typer.echo(" refresh-launcher skipped (Tauri update)") + return + _refresh_desktop_shortcuts(verbose = verbose) + + +def _release_self_exe_lock_windows() -> None: + """Rename running unsloth.exe so pip can replace it. setup.ps1 also retries.""" + if platform.system() != "Windows": + return + try: + venv_scripts = Path(sys.executable).resolve().parent + except OSError: + return + exe = venv_scripts / "unsloth.exe" + if not exe.exists(): + return + stale = exe.with_suffix(".exe.deleteme") + try: + # os.replace is atomic-overwrite on Windows; os.rename would raise + # FileExistsError if a prior aborted update left a .deleteme behind. + os.replace(exe, stale) + except OSError as e: + # Not fatal; setup.ps1 retries from a sibling process. + print(f"[update] could not rename {exe.name} -> {stale.name}: {e}") + + +def _restore_self_exe_lock_windows() -> None: + """If setup failed before pip wrote a working unsloth.exe, restore .deleteme.""" + if platform.system() != "Windows": + return + try: + venv_scripts = Path(sys.executable).resolve().parent + except OSError: + return + exe = venv_scripts / "unsloth.exe" + stale = exe.with_suffix(".exe.deleteme") + if not stale.exists(): + return + # Treat a missing or zero-byte exe as "pip didn't produce a usable + # replacement"; otherwise leave the new binary alone. + if exe.exists(): + try: + if exe.stat().st_size > 0: + return + except OSError: + return + try: + os.replace(stale, exe) + except OSError as e: + print(f"[update] could not restore {stale.name} -> {exe.name}: {e}") + + +def _cleanup_self_exe_lock_windows() -> None: + """Remove the .deleteme orphan after a successful update on Windows.""" + if platform.system() != "Windows": + return + try: + venv_scripts = Path(sys.executable).resolve().parent + except OSError: + return + stale = (venv_scripts / "unsloth.exe").with_suffix(".exe.deleteme") + try: + stale.unlink(missing_ok = True) + except OSError: + pass # ── unsloth studio reset-password ──────────────────────────────────── From dd0b5577946d0605e303fff9e329565780657ed1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:56:56 -0700 Subject: [PATCH 14/32] ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) * ci: add advisory lockfile supply-chain audit Adds a fast, focused workflow that scans every checked-in npm and cargo lockfile on PRs touching one. Default behaviour is advisory: only public indicator-of-compromise strings, versions on the public known-malicious list, and structurally broken lockfiles fail the build. Structural anomalies (missing integrity hashes, non-default registry, etc.) surface as ::warning:: annotations without gating merges, so reviewers see the audit result inline on every PR without changing the existing install behaviour. Also commits the two missing npm lockfiles the audit needs: studio/package-lock.json (Tauri CLI holder for desktop release) and studio/backend/core/data_recipe/oxc-validator/package-lock.json (oxc-parser runtime for the data-recipe validator). studio/setup.sh, studio/setup.ps1, build.sh, and pyproject.toml are intentionally left alone so the existing install path keeps working unchanged. Audit script behaviour: default mode -> exits 1 only on blocked-known-malicious, known-ioc-string, malformed-lockfile, missing-lockfile, unreadable-lockfile, or missing-toml-parser --strict -> promotes every finding to blocking (opt-in) Adds a try/except around lockfile reads so a permissions error prints a finding instead of crashing CI with a raw traceback. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test(security): update cargo regression test for advisory mode `scripts/lockfile_supply_chain_audit.py` now classifies `non-registry-cargo-source` as an advisory finding by default (returns exit 0 with a `::warning::` annotation) rather than unconditionally blocking with exit 1. Update the existing `test_malicious_cargo_lockfile_refused` to pass --strict so it keeps verifying the "refuse to install" behavior it is named for, and add a second test that pins the default-mode behavior: advisory finding emitted, exit code 0. * audit: escape Finding for GH Actions annotations `::warning::` and `::error::` workflow commands truncate the annotation message at the first newline unless the message is %-encoded per the workflow-commands spec. Since `Finding.__str__` returns three lines (kind+path, package, detail), the package and detail fields were being dropped from the GitHub Actions UI. Add a `_gha_escape()` helper that applies the spec'd escapes (`%` -> `%25`, then `\r` -> `%0D`, then `\n` -> `%0A`; the `%` replacement must happen first so the subsequent escapes are not double-encoded), wrap every Finding rendered into a workflow command with it, and pin both the helper and the end-to-end single-line emission with two new regression tests. Caught by gemini-code-assist on PR #5604. * [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/workflows/lockfile-audit.yml | 79 ++ .gitignore | 4 + scripts/lockfile_supply_chain_audit.py | 170 +++- .../oxc-validator/package-lock.json | 799 ++++++++++++++++++ studio/package-lock.json | 233 +++++ studio/package.json | 10 + .../test_lockfile_supply_chain_audit.py | 105 ++- 7 files changed, 1386 insertions(+), 14 deletions(-) create mode 100644 .github/workflows/lockfile-audit.yml create mode 100644 studio/backend/core/data_recipe/oxc-validator/package-lock.json create mode 100644 studio/package-lock.json create mode 100644 studio/package.json diff --git a/.github/workflows/lockfile-audit.yml b/.github/workflows/lockfile-audit.yml new file mode 100644 index 0000000000..9c28e21672 --- /dev/null +++ b/.github/workflows/lockfile-audit.yml @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Fast, focused supply-chain audit of every checked-in lockfile. +# +# Runs scripts/lockfile_supply_chain_audit.py on PRs that touch any +# npm or cargo lockfile, on push to main, and on a daily schedule so +# newly-published IOCs surface even when no PR opens. +# +# Default behavior is "advisory": only public indicator-of-compromise +# strings, known-malicious pinned versions, and structurally broken +# lockfiles fail the build. Structural anomalies (missing integrity, +# non-default registry, etc.) are emitted as GitHub Actions warnings +# but do not block merges. This deliberately keeps the noise floor +# low while still failing the moment a checked-in lockfile starts +# pointing at known-bad bytes. +# +# This workflow is intentionally separate from security-audit.yml: +# - security-audit.yml is the umbrella job (pip-audit + npm audit + +# cargo audit + OSV + Semgrep + secret scanning + SBOM + ...); +# it takes ~25 minutes and runs only when dep manifests change. +# - lockfile-audit.yml is a ~30 second pure-Python parse + grep on +# the lockfiles themselves; it runs on every PR that even nudges +# a lockfile so reviewers always see the audit result inline. + +name: Lockfile supply-chain audit + +on: + pull_request: + paths: + - 'studio/frontend/package-lock.json' + - 'studio/backend/core/data_recipe/oxc-validator/package-lock.json' + - 'studio/package-lock.json' + - 'studio/src-tauri/Cargo.lock' + - 'scripts/lockfile_supply_chain_audit.py' + - '.github/workflows/lockfile-audit.yml' + push: + branches: [main] + paths: + - 'studio/frontend/package-lock.json' + - 'studio/backend/core/data_recipe/oxc-validator/package-lock.json' + - 'studio/package-lock.json' + - 'studio/src-tauri/Cargo.lock' + - 'scripts/lockfile_supply_chain_audit.py' + - '.github/workflows/lockfile-audit.yml' + schedule: + - cron: '37 5 * * *' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + audit: + name: lockfile supply-chain audit + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Verify audit script parses + run: python3 -c "import ast; ast.parse(open('scripts/lockfile_supply_chain_audit.py').read())" + + - name: Run lockfile supply-chain audit + # Default mode: only known-malicious pinned versions, known IOC + # strings, and structurally broken lockfiles fail the build. + # Missing-integrity and other structural anomalies are emitted + # as ::warning:: annotations and do not gate merges. + run: python3 scripts/lockfile_supply_chain_audit.py diff --git a/.gitignore b/.gitignore index bc7d59316d..a839633790 100644 --- a/.gitignore +++ b/.gitignore @@ -229,5 +229,9 @@ log.txt setup_leo.sh server.pid *.log +# Ignore stray lockfiles; real npm projects opt back in below (npm ci needs them). package-lock.json +!studio/frontend/package-lock.json +!studio/backend/core/data_recipe/oxc-validator/package-lock.json +!studio/package-lock.json llama.cpp/ diff --git a/scripts/lockfile_supply_chain_audit.py b/scripts/lockfile_supply_chain_audit.py index ae215bf344..ffeea51c23 100644 --- a/scripts/lockfile_supply_chain_audit.py +++ b/scripts/lockfile_supply_chain_audit.py @@ -389,6 +389,18 @@ class Finding: ) +def _gha_escape(text: str) -> str: + """Escape a string for use in a GitHub Actions `::warning::` / + `::error::` workflow command message. GH Actions truncates + annotation messages at the first newline unless `\\n` is + escaped as `%0A`; carriage returns and the percent sign need + matching escapes per the workflow-commands spec. Order matters: + `%` must be replaced first so the subsequent `%0A` / `%0D` + sequences are not double-encoded. + """ + return text.replace("%", "%25").replace("\r", "%0D").replace("\n", "%0A") + + # ───────────────────────────────────────────────────────────────────── # package-lock.json audit. # ───────────────────────────────────────────────────────────────────── @@ -397,9 +409,35 @@ class Finding: def audit_npm_lockfile(path: Path) -> list[Finding]: findings: list[Finding] = [] if not path.exists(): + # A missing requested lockfile is a config error, not a clean + # audit; surface it so a deleted default cannot pass silently. + findings.append( + Finding( + path = str(path), + package = "", + kind = "missing-lockfile", + detail = ( + "expected lockfile not found; refusing to silently " + "report a clean audit for a path that was not scanned" + ), + ) + ) return findings - raw = path.read_text(encoding = "utf-8") + try: + raw = path.read_text(encoding = "utf-8") + except OSError as exc: + # Permission denied, is-a-directory, broken-pipe etc. -- surface + # as a finding instead of crashing CI with a raw traceback. + findings.append( + Finding( + path = str(path), + package = "", + kind = "unreadable-lockfile", + detail = f"could not read file: {exc}", + ) + ) + return findings try: lock = json.loads(raw) except json.JSONDecodeError as exc: @@ -553,9 +591,32 @@ _PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$") def audit_cargo_lockfile(path: Path) -> list[Finding]: findings: list[Finding] = [] if not path.exists(): + # See audit_npm_lockfile: missing lockfile is a finding. + findings.append( + Finding( + path = str(path), + package = "", + kind = "missing-lockfile", + detail = ( + "expected lockfile not found; refusing to silently " + "report a clean audit for a path that was not scanned" + ), + ) + ) return findings - raw = path.read_text(encoding = "utf-8") + try: + raw = path.read_text(encoding = "utf-8") + except OSError as exc: + findings.append( + Finding( + path = str(path), + package = "", + kind = "unreadable-lockfile", + detail = f"could not read file: {exc}", + ) + ) + return findings try: import tomllib # type: ignore[import-not-found] except ImportError: @@ -652,7 +713,31 @@ def audit_cargo_lockfile(path: Path) -> list[Finding]: # ───────────────────────────────────────────────────────────────────── -DEFAULT_NPM_LOCKFILES = ("studio/frontend/package-lock.json",) +# Finding kinds split into BLOCKING vs ADVISORY for the default run mode. +# Blocking findings come from public supply-chain attack indicators (a +# version we know is malicious, a string an attacker would have to embed +# for an attack to work). Advisory findings are structural lockfile +# anomalies (missing integrity, non-default registry, etc.) -- they +# WARN the maintainer but do not block merges. Pass --strict to make +# every finding blocking (PR-5479-style behavior for opt-in adopters). +BLOCKING_KINDS: frozenset[str] = frozenset( + { + "blocked-known-malicious", + "known-ioc-string", + # Internal-failure kinds: a structurally broken lockfile MIGHT + # be hiding a real attack, so we keep these blocking too. + "malformed-lockfile", + "missing-lockfile", + "unreadable-lockfile", + "missing-toml-parser", + } +) + +DEFAULT_NPM_LOCKFILES = ( + "studio/frontend/package-lock.json", + "studio/backend/core/data_recipe/oxc-validator/package-lock.json", + "studio/package-lock.json", +) DEFAULT_CARGO_LOCKFILES = ("studio/src-tauri/Cargo.lock",) @@ -671,7 +756,9 @@ def main(argv: list[str] | None = None) -> int: default = None, help = ( "Path to a package-lock.json (repeatable). " - "Default: studio/frontend/package-lock.json." + "Default: studio/frontend/package-lock.json, " + "studio/backend/core/data_recipe/oxc-validator/package-lock.json, " + "and studio/package-lock.json (Tauri CLI for desktop release)." ), ) parser.add_argument( @@ -683,6 +770,18 @@ def main(argv: list[str] | None = None) -> int: "Default: studio/src-tauri/Cargo.lock." ), ) + parser.add_argument( + "--strict", + action = "store_true", + help = ( + "Treat every finding as blocking (exit 1). " + "Default mode only blocks on known-malicious versions, " + "indicator-of-compromise strings, or structurally broken " + "lockfiles; everything else is printed as an advisory " + "warning with exit 0. CI should use the default; local " + "audits aiming for zero noise can opt in via --strict." + ), + ) args = parser.parse_args(argv) # SF4: require a real justification (e.g. JIRA ticket id) for the @@ -715,8 +814,15 @@ def main(argv: list[str] | None = None) -> int: return 0 root = Path(args.root).resolve() - npm_paths = [root / p for p in (args.npm_lockfile or DEFAULT_NPM_LOCKFILES)] - cargo_paths = [root / p for p in (args.cargo_lockfile or DEFAULT_CARGO_LOCKFILES)] + # Explicit --npm-lockfile/--cargo-lockfile scopes the scan to those + # paths; defaults apply only to the no-args CI invocation. + _user_explicit = args.npm_lockfile is not None or args.cargo_lockfile is not None + if _user_explicit: + npm_paths = [root / p for p in (args.npm_lockfile or ())] + cargo_paths = [root / p for p in (args.cargo_lockfile or ())] + else: + npm_paths = [root / p for p in DEFAULT_NPM_LOCKFILES] + cargo_paths = [root / p for p in DEFAULT_CARGO_LOCKFILES] all_findings: list[Finding] = [] for p in npm_paths: @@ -734,17 +840,57 @@ def main(argv: list[str] | None = None) -> int: ) return 0 + # Split findings into blocking (known-malicious / IOC / structurally + # broken) and advisory (everything else, e.g. missing integrity on a + # registry-published tarball). In default mode advisory findings are + # printed but do not change the exit code; --strict treats every + # finding as blocking. + blocking = [f for f in all_findings if f.kind in BLOCKING_KINDS] + advisory = [f for f in all_findings if f.kind not in BLOCKING_KINDS] + + if args.strict: + blocking = list(all_findings) + advisory = [] + + if advisory: + print( + f"\n[lockfile-audit] {len(advisory)} advisory finding(s) " + "(non-blocking; pass --strict to fail the build on these):\n", + file = sys.stderr, + ) + for f in advisory: + # Surface in GitHub Actions UI as a warning annotation when run + # under Actions; harmless prefix elsewhere. GH Actions + # truncates annotation messages at the first newline unless + # newlines are escaped as `%0A`, so the full multi-line + # Finding (kind + path + package + detail) only renders in + # the UI after _gha_escape collapses it onto one line. + print(f"::warning::{_gha_escape(str(f))}", file = sys.stderr) + print(file = sys.stderr) + + if not blocking: + print( + f"[lockfile-audit] OK: {len(advisory)} advisory finding(s), " + "0 blocking. Run with --strict to escalate advisory findings.", + flush = True, + ) + return 0 + print( - f"\n[lockfile-audit] FAIL: {len(all_findings)} finding(s):\n", + f"\n[lockfile-audit] FAIL: {len(blocking)} blocking finding(s):\n", file = sys.stderr, ) - for f in all_findings: - print(str(f), file = sys.stderr) + for f in blocking: + # Same %-encoding rationale as the advisory branch above: the + # GH Actions annotation is truncated at the first newline + # unless the message is escaped. + print(f"::error::{_gha_escape(str(f))}", file = sys.stderr) print(file = sys.stderr) print( - "[lockfile-audit] Refusing to proceed. Each finding above is " - "either a structural lockfile anomaly or a public indicator-of-" - "compromise. Investigate before running `npm ci` or `cargo fetch`.", + "[lockfile-audit] Refusing to proceed. Each blocking finding " + "above is either a public indicator-of-compromise, a known-" + "malicious pinned version, or a structurally broken lockfile. " + "Investigate before running `npm ci` or `cargo fetch`.", file = sys.stderr, ) return 1 diff --git a/studio/backend/core/data_recipe/oxc-validator/package-lock.json b/studio/backend/core/data_recipe/oxc-validator/package-lock.json new file mode 100644 index 0000000000..bb2ae29b23 --- /dev/null +++ b/studio/backend/core/data_recipe/oxc-validator/package-lock.json @@ -0,0 +1,799 @@ +{ + "name": "unsloth-oxc-validator-runtime", + "version": "0.0.1", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "unsloth-oxc-validator-runtime", + "version": "0.0.1", + "dependencies": { + "oxc-parser": "^0.123.0", + "oxlint": "^1.51.0" + } + }, + "node_modules/@emnapi/core": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.10.0.tgz", + "integrity": "sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==", + "license": "MIT", + "optional": true, + "peer": true, + "dependencies": { + "@emnapi/wasi-threads": "1.2.1", + "tslib": "^2.4.0" + } + }, + "node_modules/@emnapi/runtime": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.10.0.tgz", + "integrity": "sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==", + "license": "MIT", + "optional": true, + "peer": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@emnapi/wasi-threads": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz", + "integrity": "sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==", + "license": "MIT", + "optional": true, + "peer": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@napi-rs/wasm-runtime": { + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.4.tgz", + "integrity": "sha512-3NQNNgA1YSlJb/kMH1ildASP9HW7/7kYnRI2szWJaofaS1hWmbGI4H+d3+22aGzXXN9IJ+n+GiFVcGipJP18ow==", + "license": "MIT", + "optional": true, + "dependencies": { + "@tybys/wasm-util": "^0.10.1" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + }, + "peerDependencies": { + "@emnapi/core": "^1.7.1", + "@emnapi/runtime": "^1.7.1" + } + }, + "node_modules/@oxc-parser/binding-android-arm-eabi": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-android-arm-eabi/-/binding-android-arm-eabi-0.123.0.tgz", + "integrity": "sha512-EHQ58z+6DbZWokMOKg5AB1KuwrXVgfbBLuuLFfzdc7bI5A4igvdvjKMhUv1VBV+0FABiUCOjNKUmMF7ugprwbQ==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-android-arm64": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-android-arm64/-/binding-android-arm64-0.123.0.tgz", + "integrity": "sha512-BK1E0zqNoHf38nTHjnGZ+olKHSKNHh65pChjY06yhaWYP8X7yNDqhQDA4neMPRqnPBgpN4/OW1oSMrdJgDi2aw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-darwin-arm64": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-darwin-arm64/-/binding-darwin-arm64-0.123.0.tgz", + "integrity": "sha512-dkMPbtTbqU+cm+k4YGOBs4zAuq3Xu+wqjbGQvLAuVO7qHhNY4p5LBNudOmOoi0jxS8h1W6Jmlzv8MAKGpK+iDg==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-darwin-x64": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-darwin-x64/-/binding-darwin-x64-0.123.0.tgz", + "integrity": "sha512-85pic0rCd59DGdM69jI9xE/Snb2KtrfiU48QigjJXjzxUOenGvH4SAFIjFpO/2ZnI3Kz50D8pht4jKN3t2022Q==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-freebsd-x64": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-freebsd-x64/-/binding-freebsd-x64-0.123.0.tgz", + "integrity": "sha512-mjEiW6z7JtaiHMK/8aJic1lfjkKpzFwK2XFNmm187BFbtDamjGVuKNr2TEyrFEYJyZc217wokR1wrYeZGBQo4Q==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm-gnueabihf": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-0.123.0.tgz", + "integrity": "sha512-mYxigPtGt6SZfhNZBIJfuDM92cLo8XUW08WuKxzHvcmWu6xndLqwLp99Vg4uHke1AXicQEHU3Wri2X9bHF0Vlw==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm-musleabihf": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-0.123.0.tgz", + "integrity": "sha512-ttWirDC9eUBn0R4Tzz3aeDaLrx9drPdNiLJ8MXeDBFxd6cwLfTIC27qjsdfGpn942tkVIZY3sjWAnvbwDDjX7g==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm64-gnu": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-0.123.0.tgz", + "integrity": "sha512-apAHyoMNRYT+2G98Y14caZmsr5LD9PsWpGI7nXmSwK26LGiQneCU6HvHQ+d+AX+RJ5TTWZtEb2RD7OLqAC0cYQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm64-musl": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm64-musl/-/binding-linux-arm64-musl-0.123.0.tgz", + "integrity": "sha512-3r99Qa4egjO/iXUBxTlN6Ddt1YkLifG6olzvj8gkoKEK2U/MOW7mQfXRyBmuoMgmZ7O4vk41gO3d21c6VcN3yQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-ppc64-gnu": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-0.123.0.tgz", + "integrity": "sha512-Hr/Z24kUE4pjJs346g80WDwjyJGrxiw6hExJuOiME/76ZFz68y5L11UzprRkW9FN4HxBB7tLZ/fytczV2fEsiA==", + "cpu": [ + "ppc64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-riscv64-gnu": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-0.123.0.tgz", + "integrity": "sha512-sxjbhs+8WXeuoLnZ2rBmQ96gPdq3SCmz24reIltsKLUt1EDMgdaQsr7RqwBphw3QAImkMtlPQfAWDWwZyo0xDg==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-riscv64-musl": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-0.123.0.tgz", + "integrity": "sha512-d6xHHhqldA/W+VC7v8uHs24zM69Ad3HnHQ45h+uuBhCsbZx3d0E0wL2K3uJ5mYKTR6UPMFk9VMXcHWwvg1PRZQ==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-s390x-gnu": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-0.123.0.tgz", + "integrity": "sha512-+di9A5wJQlv0VodyhADjJ2rC4geyHY+uhJDl3TFjMgYhhlgLZchi9uHD5mfiUEDWHt1x7/eU2u1ge3LLazZmFw==", + "cpu": [ + "s390x" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-x64-gnu": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-x64-gnu/-/binding-linux-x64-gnu-0.123.0.tgz", + "integrity": "sha512-sh7pw2g/u6LE1TaRRQsV9Kv9+1y+CywaaNwWWP+3bnEPk/L692oTG0hmEviUlawI8v3OGC+AhbjtAD+HXWQAkg==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-x64-musl": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-x64-musl/-/binding-linux-x64-musl-0.123.0.tgz", + "integrity": "sha512-S+LoD8PiJ639JwIqK1knIeqAyYkeCbLHtAgfapszKX0yVCaYP+aer8dJxL25de9qcDjvYWVrYCkuDZzHmOl2Xw==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-openharmony-arm64": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-openharmony-arm64/-/binding-openharmony-arm64-0.123.0.tgz", + "integrity": "sha512-/65vryK11q1I+k+7ukDlwZOxUFCLYsoZBZPGZHyet5bIP5e3D8mV3uCuvpWZ9Hoe6vUZFw/nAfCrX59MeuJPgw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "openharmony" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-wasm32-wasi": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-wasm32-wasi/-/binding-wasm32-wasi-0.123.0.tgz", + "integrity": "sha512-y4OsMGQiAbZzj2Rq0LEfvhR48rQDvbvqsl/dPdn4tdf+z3H79nZuR+lQ/+KUGjD30vpVGem138sBWHFj9UR+Vg==", + "cpu": [ + "wasm32" + ], + "license": "MIT", + "optional": true, + "dependencies": { + "@napi-rs/wasm-runtime": "^1.1.2" + }, + "engines": { + "node": ">=14.0.0" + } + }, + "node_modules/@oxc-parser/binding-win32-arm64-msvc": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-0.123.0.tgz", + "integrity": "sha512-9lBqI6AXAkjYavkdpizNU3Q51uoVYfp9FJPx19hnCEdPku1jSgzSnvgmCvhCue0GziIvIvIdWgZ41wXQ3EOoBw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-win32-ia32-msvc": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-ia32-msvc/-/binding-win32-ia32-msvc-0.123.0.tgz", + "integrity": "sha512-zJbqBHwSUB7CyvAONy9ewGtQwcQj+ylOhYGETvUPp3KIYx7lolj4Gayof7iA22SU5eMSjO5COL0c8wYhmn9agA==", + "cpu": [ + "ia32" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-win32-x64-msvc": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-x64-msvc/-/binding-win32-x64-msvc-0.123.0.tgz", + "integrity": "sha512-q7RZvglQvGo3RX5ljtcGSabu2B2c0oDU/6xC3sBMhsV5KRo0PvyxLdordbEN31NTfuZu4Sgl86C76cAURZIHWA==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-project/types": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.123.0.tgz", + "integrity": "sha512-YtECP/y8Mj1lSHiUWGSRzy/C6teUKlS87dEfuVKT09LgQbUsBW1rNg+MiJ4buGu3yuADV60gbIvo9/HplA56Ew==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/Boshen" + } + }, + "node_modules/@oxlint/binding-android-arm-eabi": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-android-arm-eabi/-/binding-android-arm-eabi-1.64.0.tgz", + "integrity": "sha512-2r6Nq3XXGLHEXKkSj8JtmJ6N4gDw431DPFOg0ZoJHlNjnG6HVMm/ksQ10m0HJ8WBvwgMe1L50UHPaYZutCRPCw==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-android-arm64": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-android-arm64/-/binding-android-arm64-1.64.0.tgz", + "integrity": "sha512-ePJMpePgg7fBv+L/hVx1xXRU5/5gd5m0obLA6hPEfLXF3GjpR8idIDbY1dhQYhyz1ms2wdTccSboo6KEd2Oxtg==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-darwin-arm64": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-darwin-arm64/-/binding-darwin-arm64-1.64.0.tgz", + "integrity": "sha512-U4DMLQd10gJLuoSTLSGbfv3bGjTlUNsScm9Dgb8wwBqmCzidf1pE1pXV4doGNxqwH3KtVng1AGTINA0NvkGLvQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-darwin-x64": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-darwin-x64/-/binding-darwin-x64-1.64.0.tgz", + "integrity": "sha512-GoRIL48QWm4/TAvjN8pB1nAG+1/uqc9EdnWT9zqHeb6wsmjZtywj8VRe5aGW47Fdb64YtLOsdLqVxOvQuz98Wg==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-freebsd-x64": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-freebsd-x64/-/binding-freebsd-x64-1.64.0.tgz", + "integrity": "sha512-5dFkv4tkg7PxJJGS9/OjrJwjhuHczrd3OQOkRE0wHcLM+ncUnULtzEPWjqGOxTXxZnLWcB91bGiIznx89TVXyQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm-gnueabihf": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.64.0.tgz", + "integrity": "sha512-jsBqMLl/uOL5+Kq/+BtK9FrmiNGUbx8SiyZXv+WlUxA45KuwcLu9BfiSIL3I3DBDgWM3yZizDITnTK9BcqNBQg==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm-musleabihf": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-1.64.0.tgz", + "integrity": "sha512-1lrj8At/Uuc9GhjrVFBQo0NEjfBrTkzpmtHIGAhNnIXqn1CAyGL+qrztUsXb2GIluJrpl9Q7qRLJOb/NqydacQ==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm64-gnu": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.64.0.tgz", + "integrity": "sha512-HpSQbubwh03mMhAdy2BYtad/fsY8vDFHDAb6bUwuCYg2VD3xCQgn6ArKcO0oZyLCheacKTv4PrF3Mfu5hgoE2g==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm64-musl": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.64.0.tgz", + "integrity": "sha512-00QQ0h0Y7u0G69BgiH3+ky2aaq/QvkDL6DYok8htIuJHxybiux5aQ8jwmg8qIk9wha6UagUP2BAwAzbemcJbpg==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-ppc64-gnu": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.64.0.tgz", + "integrity": "sha512-2GaimTV6EMW+s5HS0An3oGbQme3BgHswvfVdGk3EB57Xe9+/gyT+Qd7lNVzb3rtir52vbIPzXfaYArzs5b5zcw==", + "cpu": [ + "ppc64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-riscv64-gnu": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-1.64.0.tgz", + "integrity": "sha512-H46AtFb9wypjoVwGdlxrm0DsD809NGmtiK9HiyPKTxkSte2YjhC4S+00rOIrwCaxcyPiGid3Y3OMXp5KMAkGZw==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-riscv64-musl": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-1.64.0.tgz", + "integrity": "sha512-HEgsidjjvvyzdg82icYkuFCf7REDV7B9JFwbIMbVwrKLBY0MrXX+bku3POn/hduZ2yW91IyVDUMq0Bf02KwXQw==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-s390x-gnu": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.64.0.tgz", + "integrity": "sha512-Axvm8qryotmKN00P5w4JapaSjvP2LOSbdbBJiX+2SuHd3QzhW7TUc8skqgw+ahQZ5DmzEYeHCqauvW8f32Ns6Q==", + "cpu": [ + "s390x" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-x64-gnu": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.64.0.tgz", + "integrity": "sha512-cR60vSd7+m+KRZ3GQGfDxWwahW5RMXg0qlGvAluZr0fTUYvw0H9N9AXAF/M/PMqgytyqvVNmBAkJG9l7U30Y1g==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-x64-musl": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-x64-musl/-/binding-linux-x64-musl-1.64.0.tgz", + "integrity": "sha512-2u/aPZ9pEg7HnvZPDsHxUGNnrpr4qaHi+mCgLgpt+LYRzPrS4Px4wPfkIdRdr2GvKnaYyt+XSlto0Vm5sbStTg==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-openharmony-arm64": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-openharmony-arm64/-/binding-openharmony-arm64-1.64.0.tgz", + "integrity": "sha512-kfhkGfCdoXLSxEkrhDlJrvBYajGmq+ma4EMc53dsOWTq+rIBOlI0vTBmpZNnM5oH2LY/K/w1HAK+UQEgjgpVUg==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "openharmony" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-win32-arm64-msvc": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.64.0.tgz", + "integrity": "sha512-r/cNKBFieONoVu2bb1KkVouq9W+edDUgHumXJGphCRRj+U0xaD4nanrw8ZOqo0IsutPkEM4vCcGBpak6x5aXMg==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-win32-ia32-msvc": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-win32-ia32-msvc/-/binding-win32-ia32-msvc-1.64.0.tgz", + "integrity": "sha512-tUw0xUUwEFVZbpJoeCblkv8SJA4Xz3CdXCJbAnBsiNLyxDrk2tLcxEAS6M73Q7hHHDg3OtwI8vZVK3t5RJt4Gw==", + "cpu": [ + "ia32" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-win32-x64-msvc": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.64.0.tgz", + "integrity": "sha512-9CBR+LO0JVST87fNTzzNxS5I29jIUO5gxT9i9+M3SDHHALElj9sY1Prf12tad3vIRC6OD7Ehtvvh+sn13vSwHw==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@tybys/wasm-util": { + "version": "0.10.2", + "resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.2.tgz", + "integrity": "sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg==", + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/oxc-parser": { + "version": "0.123.0", + "resolved": "https://registry.npmjs.org/oxc-parser/-/oxc-parser-0.123.0.tgz", + "integrity": "sha512-F6ak0tFc01ZGbl5KxvLDQ2K005Z086mp3ByCQBDhUjqXLkapGUkMuJSsYixncdEpkLlcRDcruHR71LD339ADUA==", + "license": "MIT", + "dependencies": { + "@oxc-project/types": "^0.123.0" + }, + "engines": { + "node": "^20.19.0 || >=22.12.0" + }, + "funding": { + "url": "https://github.com/sponsors/Boshen" + }, + "optionalDependencies": { + "@oxc-parser/binding-android-arm-eabi": "0.123.0", + "@oxc-parser/binding-android-arm64": "0.123.0", + "@oxc-parser/binding-darwin-arm64": "0.123.0", + "@oxc-parser/binding-darwin-x64": "0.123.0", + "@oxc-parser/binding-freebsd-x64": "0.123.0", + "@oxc-parser/binding-linux-arm-gnueabihf": "0.123.0", + "@oxc-parser/binding-linux-arm-musleabihf": "0.123.0", + "@oxc-parser/binding-linux-arm64-gnu": "0.123.0", + "@oxc-parser/binding-linux-arm64-musl": "0.123.0", + "@oxc-parser/binding-linux-ppc64-gnu": "0.123.0", + "@oxc-parser/binding-linux-riscv64-gnu": "0.123.0", + "@oxc-parser/binding-linux-riscv64-musl": "0.123.0", + "@oxc-parser/binding-linux-s390x-gnu": "0.123.0", + "@oxc-parser/binding-linux-x64-gnu": "0.123.0", + "@oxc-parser/binding-linux-x64-musl": "0.123.0", + "@oxc-parser/binding-openharmony-arm64": "0.123.0", + "@oxc-parser/binding-wasm32-wasi": "0.123.0", + "@oxc-parser/binding-win32-arm64-msvc": "0.123.0", + "@oxc-parser/binding-win32-ia32-msvc": "0.123.0", + "@oxc-parser/binding-win32-x64-msvc": "0.123.0" + } + }, + "node_modules/oxlint": { + "version": "1.64.0", + "resolved": "https://registry.npmjs.org/oxlint/-/oxlint-1.64.0.tgz", + "integrity": "sha512-Star3SNpWPeWFPw7kRXIhXUSn6fdiAl25q15CQzH/9WaOtG6e9CWTc25vNZOCr4PE1yEP1GtKJKIKglhj3OmEQ==", + "license": "MIT", + "bin": { + "oxlint": "bin/oxlint" + }, + "engines": { + "node": "^20.19.0 || >=22.12.0" + }, + "funding": { + "url": "https://github.com/sponsors/Boshen" + }, + "optionalDependencies": { + "@oxlint/binding-android-arm-eabi": "1.64.0", + "@oxlint/binding-android-arm64": "1.64.0", + "@oxlint/binding-darwin-arm64": "1.64.0", + "@oxlint/binding-darwin-x64": "1.64.0", + "@oxlint/binding-freebsd-x64": "1.64.0", + "@oxlint/binding-linux-arm-gnueabihf": "1.64.0", + "@oxlint/binding-linux-arm-musleabihf": "1.64.0", + "@oxlint/binding-linux-arm64-gnu": "1.64.0", + "@oxlint/binding-linux-arm64-musl": "1.64.0", + "@oxlint/binding-linux-ppc64-gnu": "1.64.0", + "@oxlint/binding-linux-riscv64-gnu": "1.64.0", + "@oxlint/binding-linux-riscv64-musl": "1.64.0", + "@oxlint/binding-linux-s390x-gnu": "1.64.0", + "@oxlint/binding-linux-x64-gnu": "1.64.0", + "@oxlint/binding-linux-x64-musl": "1.64.0", + "@oxlint/binding-openharmony-arm64": "1.64.0", + "@oxlint/binding-win32-arm64-msvc": "1.64.0", + "@oxlint/binding-win32-ia32-msvc": "1.64.0", + "@oxlint/binding-win32-x64-msvc": "1.64.0" + }, + "peerDependencies": { + "oxlint-tsgolint": ">=0.22.1" + }, + "peerDependenciesMeta": { + "oxlint-tsgolint": { + "optional": true + } + } + }, + "node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", + "license": "0BSD", + "optional": true + } + } +} diff --git a/studio/package-lock.json b/studio/package-lock.json new file mode 100644 index 0000000000..b5ed4ea2f3 --- /dev/null +++ b/studio/package-lock.json @@ -0,0 +1,233 @@ +{ + "name": "unsloth-studio-tauri-cli", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "unsloth-studio-tauri-cli", + "version": "0.0.0", + "license": "AGPL-3.0-only", + "devDependencies": { + "@tauri-apps/cli": "2.10.1" + } + }, + "node_modules/@tauri-apps/cli": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli/-/cli-2.10.1.tgz", + "integrity": "sha512-jQNGF/5quwORdZSSLtTluyKQ+o6SMa/AUICfhf4egCGFdMHqWssApVgYSbg+jmrZoc8e1DscNvjTnXtlHLS11g==", + "dev": true, + "license": "Apache-2.0 OR MIT", + "bin": { + "tauri": "tauri.js" + }, + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/tauri" + }, + "optionalDependencies": { + "@tauri-apps/cli-darwin-arm64": "2.10.1", + "@tauri-apps/cli-darwin-x64": "2.10.1", + "@tauri-apps/cli-linux-arm-gnueabihf": "2.10.1", + "@tauri-apps/cli-linux-arm64-gnu": "2.10.1", + "@tauri-apps/cli-linux-arm64-musl": "2.10.1", + "@tauri-apps/cli-linux-riscv64-gnu": "2.10.1", + "@tauri-apps/cli-linux-x64-gnu": "2.10.1", + "@tauri-apps/cli-linux-x64-musl": "2.10.1", + "@tauri-apps/cli-win32-arm64-msvc": "2.10.1", + "@tauri-apps/cli-win32-ia32-msvc": "2.10.1", + "@tauri-apps/cli-win32-x64-msvc": "2.10.1" + } + }, + "node_modules/@tauri-apps/cli-darwin-arm64": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.10.1.tgz", + "integrity": "sha512-Z2OjCXiZ+fbYZy7PmP3WRnOpM9+Fy+oonKDEmUE6MwN4IGaYqgceTjwHucc/kEEYZos5GICve35f7ZiizgqEnQ==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-darwin-x64": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-x64/-/cli-darwin-x64-2.10.1.tgz", + "integrity": "sha512-V/irQVvjPMGOTQqNj55PnQPVuH4VJP8vZCN7ajnj+ZS8Kom1tEM2hR3qbbIRoS3dBKs5mbG8yg1WC+97dq17Pw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-linux-arm-gnueabihf": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm-gnueabihf/-/cli-linux-arm-gnueabihf-2.10.1.tgz", + "integrity": "sha512-Hyzwsb4VnCWKGfTw+wSt15Z2pLw2f0JdFBfq2vHBOBhvg7oi6uhKiF87hmbXOBXUZaGkyRDkCHsdzJcIfoJC2w==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-linux-arm64-gnu": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-gnu/-/cli-linux-arm64-gnu-2.10.1.tgz", + "integrity": "sha512-OyOYs2t5GkBIvyWjA1+h4CZxTcdz1OZPCWAPz5DYEfB0cnWHERTnQ/SLayQzncrT0kwRoSfSz9KxenkyJoTelA==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-linux-arm64-musl": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-musl/-/cli-linux-arm64-musl-2.10.1.tgz", + "integrity": "sha512-MIj78PDDGjkg3NqGptDOGgfXks7SYJwhiMh8SBoZS+vfdz7yP5jN18bNaLnDhsVIPARcAhE1TlsZe/8Yxo2zqg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-linux-riscv64-gnu": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-riscv64-gnu/-/cli-linux-riscv64-gnu-2.10.1.tgz", + "integrity": "sha512-X0lvOVUg8PCVaoEtEAnpxmnkwlE1gcMDTqfhbefICKDnOTJ5Est3qL0SrWxizDackIOKBcvtpejrSiVpuJI1kw==", + "cpu": [ + "riscv64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-linux-x64-gnu": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.10.1.tgz", + "integrity": "sha512-2/12bEzsJS9fAKybxgicCDFxYD1WEI9kO+tlDwX5znWG2GwMBaiWcmhGlZ8fi+DMe9CXlcVarMTYc0L3REIRxw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-linux-x64-musl": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.10.1.tgz", + "integrity": "sha512-Y8J0ZzswPz50UcGOFuXGEMrxbjwKSPgXftx5qnkuMs2rmwQB5ssvLb6tn54wDSYxe7S6vlLob9vt0VKuNOaCIQ==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-win32-arm64-msvc": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-arm64-msvc/-/cli-win32-arm64-msvc-2.10.1.tgz", + "integrity": "sha512-iSt5B86jHYAPJa/IlYw++SXtFPGnWtFJriHn7X0NFBVunF6zu9+/zOn8OgqIWSl8RgzhLGXQEEtGBdR4wzpVgg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-win32-ia32-msvc": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-ia32-msvc/-/cli-win32-ia32-msvc-2.10.1.tgz", + "integrity": "sha512-gXyxgEzsFegmnWywYU5pEBURkcFN/Oo45EAwvZrHMh+zUSEAvO5E8TXsgPADYm31d1u7OQU3O3HsYfVBf2moHw==", + "cpu": [ + "ia32" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@tauri-apps/cli-win32-x64-msvc": { + "version": "2.10.1", + "resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.10.1.tgz", + "integrity": "sha512-6Cn7YpPFwzChy0ERz6djKEmUehWrYlM+xTaNzGPgZocw3BD7OfwfWHKVWxXzdjEW2KfKkHddfdxK1XXTYqBRLg==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0 OR MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } + } + } +} diff --git a/studio/package.json b/studio/package.json new file mode 100644 index 0000000000..ac4789e874 --- /dev/null +++ b/studio/package.json @@ -0,0 +1,10 @@ +{ + "name": "unsloth-studio-tauri-cli", + "version": "0.0.0", + "private": true, + "description": "Lockfile holder for @tauri-apps/cli used by the desktop release workflow. Not a real npm package; `npm ci --prefix studio` resolves the pinned Tauri CLI from this directory's package-lock.json.", + "license": "AGPL-3.0-only", + "devDependencies": { + "@tauri-apps/cli": "2.10.1" + } +} diff --git a/tests/security/test_lockfile_supply_chain_audit.py b/tests/security/test_lockfile_supply_chain_audit.py index 275794da9c..483bb9e763 100644 --- a/tests/security/test_lockfile_supply_chain_audit.py +++ b/tests/security/test_lockfile_supply_chain_audit.py @@ -27,9 +27,12 @@ def _run_auditor( root: Path, npm_lockfiles: list[Path] | None = None, cargo_lockfiles: list[Path] | None = None, + strict: bool = False, timeout: int = 30, ) -> subprocess.CompletedProcess: cmd = [sys.executable, str(SCRIPT), "--root", str(root)] + if strict: + cmd.append("--strict") for p in npm_lockfiles or []: cmd.extend(["--npm-lockfile", str(p)]) for p in cargo_lockfiles or []: @@ -170,6 +173,32 @@ checksum = "0000000000000000000000000000000000000000000000000000000000000000" def test_malicious_cargo_lockfile_refused(tmp_path): """Inline Cargo.lock with `source = "git+https://example.com/..."` must trip the `non-registry-cargo-source` check. + + `non-registry-cargo-source` is an advisory finding kind in the + auditor's default mode (per the audit script's BLOCKING_KINDS + set). To exercise the historical "refuse to install" behavior we + pass --strict here; that promotes every finding to blocking and + keeps the test honest about its intent (detection + refusal). + """ + lockfile = tmp_path / "Cargo.lock" + lockfile.write_text(_MALICIOUS_CARGO_LOCK) + proc = _run_auditor( + root = tmp_path, + npm_lockfiles = [FIXTURES / "clean_lockfile.json"], + cargo_lockfiles = [lockfile], + strict = True, + ) + assert proc.returncode == 1 + combined = proc.stdout + proc.stderr + assert "non-registry-cargo-source" in combined + assert "git+https://example.com" in combined + + +def test_malicious_cargo_lockfile_default_mode_advisory(tmp_path): + """Default (non-strict) mode classifies `non-registry-cargo-source` + as advisory: the finding is still emitted as a `::warning::` + annotation but the process exits 0 so the build is not gated. + Regression test for the advisory/strict split. """ lockfile = tmp_path / "Cargo.lock" lockfile.write_text(_MALICIOUS_CARGO_LOCK) @@ -178,10 +207,13 @@ def test_malicious_cargo_lockfile_refused(tmp_path): npm_lockfiles = [FIXTURES / "clean_lockfile.json"], cargo_lockfiles = [lockfile], ) - assert proc.returncode == 1 + assert proc.returncode == 0, ( + f"expected exit 0 (advisory), got {proc.returncode}\n" + f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}" + ) combined = proc.stdout + proc.stderr assert "non-registry-cargo-source" in combined - assert "git+https://example.com" in combined + assert "advisory finding" in combined def test_audit_cargo_lockfile_direct_call(tmp_path): @@ -192,6 +224,75 @@ def test_audit_cargo_lockfile_direct_call(tmp_path): assert "non-registry-cargo-source" in kinds +# --------------------------------------------------------------------------- +# GitHub Actions annotation escape: ::warning:: / ::error:: messages +# are truncated at the first newline unless escaped, so the multi-line +# Finding must be collapsed via the spec'd %0A / %0D / %25 encoding. +# --------------------------------------------------------------------------- + + +def test_gha_escape_collapses_finding_to_one_line(): + """`_gha_escape()` must collapse newlines (`%0A`), carriage + returns (`%0D`), and percent signs (`%25`) so that + `::warning::` / `::error::` render the full finding + in the GitHub Actions UI annotation instead of being truncated + at the first newline. The `%` replacement must happen first or + the subsequent `%0A` / `%0D` escapes get double-encoded. + """ + assert lsa._gha_escape("a\nb\nc") == "a%0Ab%0Ac" + assert lsa._gha_escape("a\rb") == "a%0Db" + assert lsa._gha_escape("100%") == "100%25" + # Order regression: `%` must escape before `\n` so the literal + # text `a%b\nc` becomes `a%25b%0Ac`, not `a%250Ab%0Ac`. + assert lsa._gha_escape("a%b\nc") == "a%25b%0Ac" + + f = lsa.Finding( + path = "/x/lock.json", + package = "node_modules/foo", + kind = "missing-integrity-hash", + detail = "bad stuff", + ) + escaped = lsa._gha_escape(str(f)) + assert "\n" not in escaped + assert "%0A" in escaped + assert "missing-integrity-hash" in escaped + assert "node_modules/foo" in escaped + assert "bad stuff" in escaped + + +def test_advisory_finding_emitted_as_single_line_annotation(tmp_path): + """End-to-end check: the `::warning::` line emitted for an + advisory finding must be a SINGLE physical line (the rest of + the Finding is `%0A`-escaped inside the message). Regression + test for the gemini-code-assist review on PR #5604: without + `_gha_escape`, GitHub Actions truncates the annotation after + `[kind] path` and the package + detail fields never render. + """ + lockfile = tmp_path / "Cargo.lock" + lockfile.write_text(_MALICIOUS_CARGO_LOCK) + proc = _run_auditor( + root = tmp_path, + npm_lockfiles = [FIXTURES / "clean_lockfile.json"], + cargo_lockfiles = [lockfile], + ) + warning_lines = [ + line for line in proc.stderr.splitlines() if line.startswith("::warning::") + ] + assert warning_lines, ( + "expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}" + ) + for line in warning_lines: + # Single physical line: kind, package, detail all present + # via %0A escape, not split across stderr lines. + assert "%0A" in line, ( + f"::warning:: line has no %0A escape; multi-line text " + f"would be truncated by GH Actions:\n{line}" + ) + assert "non-registry-cargo-source" in line + assert "package:" in line + assert "detail:" in line + + # --------------------------------------------------------------------------- # SF4: skip env var requires a justification value. # --------------------------------------------------------------------------- From 06526f9d6a4dc0ecf2c325eb4d8131eca9e1dc2b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:58:32 -0700 Subject: [PATCH 15/32] loader: import FORCE_FLOAT32 from unsloth_zoo (single source of truth) (#5610) * loader: import FORCE_FLOAT32 from unsloth_zoo (single source of truth) unsloth_zoo now owns the FORCE_FLOAT32 list in unsloth_zoo/model_lists.py (re-exported as the top-level `unsloth_zoo.FORCE_FLOAT32`). The CUDA loader here imports from there so the MLX loader (unsloth_zoo.mlx.loader) and the CUDA loader stay in sync from a single edit, and the bf16->fp16 downcast warning added in unsloth-zoo PR #670 gates on the same list. Companion to unsloth-zoo PR #670. * loader: add inline FORCE_FLOAT32 fallback for old unsloth_zoo installs If a user upgrades unsloth without upgrading unsloth_zoo, the previously unconditional `from unsloth_zoo import FORCE_FLOAT32` would raise ImportError at module import time, killing the whole package. Wrap the import in try/except and fall back to an inline list that mirrors unsloth_zoo.model_lists.FORCE_FLOAT32 byte-for-byte, so the module loads cleanly on any zoo version while still preferring zoo as the single source of truth when present. --- unsloth/models/loader.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index c10443e289..a7dcaa4d88 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -99,15 +99,22 @@ from ._utils import ( fast_inference_setup, ) -global FORCE_FLOAT32 -# Forces float32 precision since float16 goes to infinity -FORCE_FLOAT32 = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n - "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) - "gemma3n", - "gpt_oss", - "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training -] +# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers +# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working. +# Fallback list mirrors zoo for users who upgrade unsloth without upgrading +# unsloth_zoo (so this module never fails at import). +try: + from unsloth_zoo import FORCE_FLOAT32 # noqa: F401 +except ImportError: + global FORCE_FLOAT32 + # Forces float32 precision since float16 goes to infinity + FORCE_FLOAT32 = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) + "gemma3n", + "gpt_oss", + "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training + ] global DISABLE_COMPILE_MODEL_NAMES # Must be alphabetically sorted for each entry @@ -1381,7 +1388,6 @@ class FastModel(FastBaseModel): for model_type_arch in model_types: if model_type_arch != "siglip": break - global FORCE_FLOAT32 for disable_name in FORCE_FLOAT32: # add comma to model_types_all matching in case of exact match for end if ( From b7acc14d0ce2d4601be9ca34bffe6f791c5bb63c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 05:58:44 -0700 Subject: [PATCH 16/32] fix(peft): expose finetune_last_n_layers for parity with mlx-lm CLI (#5564) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(peft): expose finetune_last_n_layers for parity with mlx-lm CLI mlx-lm's lora CLI defaults `CONFIG_DEFAULTS['num_layers']=16` (mlx_lm/lora.py:56), so it applies LoRA only to the LAST 16 transformer blocks. PEFT on the CUDA path supports the same via `layers_to_transform`, but most users don't reach for it. This commit adds a `finetune_last_n_layers` convenience parameter to both `FastLlamaModel.get_peft_model` and `FastBaseModel.get_peft_model` (vision/multi-modal). When set, it fills `layers_to_transform` automatically with the last N blocks, mirroring mlx-lm CLI's behavior AND `unsloth_zoo.mlx.loader.FastMLXModel.get_peft_model`. A single config value now controls layer-selection consistently across CUDA, MLX (zoo), and mlx-lm CLI paths. Default is None (= train all layers, current behavior unchanged). When set, the value is clamped to [1, total_transformer_layers] so callers can't accidentally over- or under-select. The total is read from `config.num_hidden_layers` (or aliases), falling through to `config.text_config.num_hidden_layers` for VLMs. Why this matters: with the same fixture/seed, training the last N layers vs all layers picks a different basin under stochastic LoRA init. Empirically (n=15 seeds, gemma-3-270m-it single-row LoRA memorization, MLX path) last-16 hits 67% greedy-decode pass rate vs all-18 at 47%. The teacher-forced completion loss is 0 in both — the model memorizes either way; only the first- token argmax distribution differs. CUDA fp32 shows the same pattern. Aligning the layer selection puts CUDA + MLX + mlx-lm all in the same basin family for parity comparisons. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * peft: trim verbose finetune_last_n_layers comments Per code-comment policy: parameter name is self-documenting, the clamp and range() construction are obvious. Rationale (mlx-lm CLI parity, empirical pass-rate data) lives in commit 106c1df4's message and the PR description. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_finetune_last_n_layers.py | 92 ++++++++++++++++++++++++++++ unsloth/models/llama.py | 10 +++ unsloth/models/vision.py | 36 +++++++++++ 3 files changed, 138 insertions(+) create mode 100644 tests/test_finetune_last_n_layers.py diff --git a/tests/test_finetune_last_n_layers.py b/tests/test_finetune_last_n_layers.py new file mode 100644 index 0000000000..fa0e2af5bc --- /dev/null +++ b/tests/test_finetune_last_n_layers.py @@ -0,0 +1,92 @@ +# Unsloth - 2x faster, 70% less memory LLM finetuning +# Tests for the `finetune_last_n_layers` parity knob (CUDA side). +# +# Mirrors unsloth-zoo's `FastMLXModel.get_peft_model` parameter. +# mlx-lm CLI's CONFIG_DEFAULTS['num_layers']=16 applies LoRA to the +# last 16 transformer blocks only. On the CUDA path, PEFT exposes +# `layers_to_transform` to do the same. This convenience knob fills +# `layers_to_transform` for the user when set, matching mlx-lm CLI +# AND unsloth-zoo's MLX path with a single config value. +# +# The tests intentionally avoid pulling in CUDA / a real model +# checkpoint — they exercise only the helper that translates +# `finetune_last_n_layers` into `layers_to_transform`. + +from __future__ import annotations + +import pytest + + +def test_get_total_transformer_layers_reads_num_hidden_layers(): + from unsloth.models.vision import _get_total_transformer_layers + + class FakeConfig: + num_hidden_layers = 18 + + class FakeModel: + config = FakeConfig() + + assert _get_total_transformer_layers(FakeModel()) == 18 + + +def test_get_total_transformer_layers_reads_text_config(): + from unsloth.models.vision import _get_total_transformer_layers + + class TextConfig: + num_hidden_layers = 24 + + class FakeConfig: + text_config = TextConfig() + + class FakeModel: + config = FakeConfig() + + # No num_hidden_layers at top level — should fall through to text_config. + assert _get_total_transformer_layers(FakeModel()) == 24 + + +def test_get_total_transformer_layers_handles_alternative_attr_names(): + from unsloth.models.vision import _get_total_transformer_layers + + for attr in ("n_layer", "n_layers", "num_layers"): + cfg = type("Cfg", (), {attr: 12})() + model = type("M", (), {"config": cfg})() + assert _get_total_transformer_layers(model) == 12 + + +def test_get_total_transformer_layers_returns_none_when_unknown(): + from unsloth.models.vision import _get_total_transformer_layers + + class FakeConfig: + pass + + class FakeModel: + config = FakeConfig() + + assert _get_total_transformer_layers(FakeModel()) is None + + +def test_get_total_transformer_layers_returns_none_for_missing_config(): + from unsloth.models.vision import _get_total_transformer_layers + + class FakeModel: + pass + + assert _get_total_transformer_layers(FakeModel()) is None + + +def test_finetune_last_n_layers_signature_present_on_llama_and_vision(): + """Both entry points must expose the new parameter with default None.""" + import inspect + from unsloth.models.llama import FastLlamaModel + from unsloth.models.vision import FastBaseModel + + for cls in (FastLlamaModel, FastBaseModel): + sig = inspect.signature(cls.get_peft_model) + assert ( + "finetune_last_n_layers" in sig.parameters + ), f"{cls.__name__}.get_peft_model missing finetune_last_n_layers" + assert sig.parameters["finetune_last_n_layers"].default is None, ( + f"{cls.__name__}.get_peft_model: finetune_last_n_layers default " + f"must be None to preserve historical behavior" + ) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6ddfe04d21..20b515c711 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2831,6 +2831,7 @@ class FastLlamaModel: bias = "none", layers_to_transform = None, layers_pattern = None, + finetune_last_n_layers = None, use_gradient_checkpointing = "unsloth", random_state = 3407, max_seq_length = 2048, # not used anymore @@ -2863,6 +2864,7 @@ class FastLlamaModel: bias = bias, layers_to_transform = layers_to_transform, layers_pattern = layers_pattern, + finetune_last_n_layers = finetune_last_n_layers, use_gradient_checkpointing = use_gradient_checkpointing, random_state = random_state, max_seq_length = max_seq_length, @@ -3160,6 +3162,14 @@ class FastLlamaModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + if finetune_last_n_layers is not None and layers_to_transform is None: + from .vision import _get_total_transformer_layers + + _total_layers = _get_total_transformer_layers(model) + if _total_layers is not None and _total_layers > 0: + _n = max(1, min(int(finetune_last_n_layers), _total_layers)) + layers_to_transform = list(range(_total_layers - _n, _total_layers)) + arguments = dict( r = r, lora_alpha = lora_alpha, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index df371e00c8..73ef1db7e3 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -547,6 +547,35 @@ def _construct_vlm_processor_fallback( return None +def _get_total_transformer_layers(model): + """Best-effort total transformer block count across HF model shapes. + Returns None if not determinable; caller should skip the conversion.""" + cfg = getattr(model, "config", None) + if cfg is None: + return None + for name in ( + "num_hidden_layers", + "n_layer", + "n_layers", + "num_layers", + ): + v = getattr(cfg, name, None) + if isinstance(v, int) and v > 0: + return v + text_cfg = getattr(cfg, "text_config", None) + if text_cfg is not None: + for name in ( + "num_hidden_layers", + "n_layer", + "n_layers", + "num_layers", + ): + v = getattr(text_cfg, name, None) + if isinstance(v, int) and v > 0: + return v + return None + + class FastBaseModel: @staticmethod def from_pretrained( @@ -1319,6 +1348,7 @@ class FastBaseModel: finetune_language_layers = True, finetune_attention_modules = True, finetune_mlp_modules = True, + finetune_last_n_layers = None, layers_to_transform = None, layers_pattern = None, use_gradient_checkpointing = "unsloth", @@ -1417,6 +1447,12 @@ class FastBaseModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + if finetune_last_n_layers is not None and layers_to_transform is None: + _total_layers = _get_total_transformer_layers(model) + if _total_layers is not None and _total_layers > 0: + n = max(1, min(int(finetune_last_n_layers), _total_layers)) + layers_to_transform = list(range(_total_layers - n, _total_layers)) + # Get only allowed parameters for LoraConfig local_variables = { **locals(), From 27d4aced5928b4475678b5aa5dbce26cbca15602 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:17:04 -0700 Subject: [PATCH 17/32] studio: add --spec-draft-n-max toggle for MTP speculative decoding (#5582) * studio: add --spec-draft-n-max toggle for MTP speculative decoding Surface llama-server's --spec-draft-n-max as a first-class LoadRequest field so users can tune the MTP draft tree size from the chat settings panel. Default behaviour is unchanged: when the caller omits spec_draft_n_max, the existing platform defaults still apply (6 on GPU, 3 on CPU/Mac). Why this matters: on context-constrained loads the draft KV cache competes with the target model's KV cache for VRAM. Lowering spec_draft_n_max reduces that pressure, lets a larger user context fit, and recovers throughput; raising it pays off when draft acceptance is high enough to amortise the extra cache. Backend - LoadRequest gains an optional spec_draft_n_max: int (1..16). - LlamaCppBackend.load_model accepts and persists the override on self._spec_draft_n_max, used in place of the hardcoded 6/3 in the MTP emit branch. - LoadResponse and InferenceStatusResponse echo the active value (None when the platform default is in effect) so the UI can hydrate the input on refresh. - _already_in_target_state and _request_matches_loaded_settings compare spec_draft_n_max alongside speculative_type so a value change triggers a reload rather than no-op'ing. - strip_shadowing_flags now strips inherited --spec-* extras when either speculative_type or spec_draft_n_max is in fields_set, so an inherited --spec-draft-n-max cannot last-wins-override a fresh request's first-class field. Frontend - LoadModelRequest, LoadModelResponse, InferenceStatusResponse TypeScript shapes get spec_draft_n_max. - chat-runtime-store gains specDraftNMax / loadedSpecDraftNMax and a setter, hydrated from /v1/status and /v1/load. - chat-settings-sheet renders a "Draft Tokens" numeric input directly under the Speculative Decoding switch when that switch is on. Toggling the switch off clears the override; the Reset button restores the loaded value. Tests - Four new regression tests cover _already_in_target_state with matching / mismatching / non-MTP / unset spec_draft_n_max. - Existing test_llama_server_args.py and test_llama_cpp_mtp_detection.py green: 141 passed locally. * studio: add --spec-draft-p-min and --spec-draft-p-split to spec strip set llama.cpp server documents --spec-draft-p-min (default 0.75, min draft acceptance probability) and --spec-draft-p-split (default 0.10). Both are first-class spec-decoding knobs that should travel with the rest of the --spec-* family when an Apply re-sets speculative_type, so an inherited override doesn't leak across a fresh load. * studio/tests: skip MTP capability-probe tests on Windows The four probe_server_capabilities tests use a bash stub written to tmp_path/llama-server, which Windows' subprocess can't execute directly (no shebang resolution, .bat / .cmd would be needed). Mark them skipif sys.platform == 'win32' so the rest of the MTP plumbing suite stays green on Windows CI. Unix coverage is unchanged. * studio: lower MTP GPU default --spec-draft-n-max from 6 to 2 Bench on B200 / Qwen3.6-27B-MTP-GGUF UD-Q4_K_XL across five prompt types (essay, code, story, math, science) with greedy temp=0: prompt OFF n=1 n=2 n=3 n=6 essay 79.1 93.4 93.8 84.7 64.6 code 79.1 104.4 116.6 113.5 103.0 story 79.1 99.2 105.7 101.8 88.9 math 79.1 100.8 110.8 111.8 98.2 science 79.1 100.1 110.8 110.8 102.9 The previous hardcoded GPU default of 6 was 17% SLOWER than spec-off on the essay prompt (64.6 vs 79.1 t/s) and 11-50% slower than n=2 on the rest. n=2 wins on 4/5 prompts with a 1.18x-1.47x speedup vs OFF; n=3 wins on the math prompt by a hair. n=6 collapses once acceptance rate drops past n=3 -- wasted draft decode dominates the per-step budget. Matches the dataset README ("n_max=2 is the sweet spot for 36 of 42 quants"). Keeps CPU/Mac default at 3, which empirically tracks the narrower ngram+MTP chained budget on those platforms. Users who want the old behaviour can pass spec_draft_n_max in LoadRequest (the toggle this PR also adds) or --spec-draft-n-max via llama_extra_args. * studio: skip MTP auto-promote on sub-2B models, backfill chat usage Two MTP-visibility fixes uncovered while bisecting llama.cpp post-#22673 on Qwen3.6-27B-MTP-GGUF UD-Q4_K_XL on B200. Size gate. Direct llama-server bench (no Studio measurement loop) at n_predict=192 across 9 prompts shows MTP regresses vs spec-off on sub-2B dense models because draft cost exceeds savings: Qwen3.5-0.8B Q4_K_XL GPU: 452.0 OFF -> 283.4 t/s n=2 (0.63x) CPU: 84.5 OFF -> 64.9 t/s n=3 (0.77x) Qwen3.5-4B Q4_K_XL GPU: 241.0 OFF -> 258.2 t/s n=2 (1.07x) Qwen3.5-9B Q4_K_XL GPU: 201.6 OFF -> 228.9 t/s n=2 (1.14x) Qwen3.5-27B Q4_K_XL GPU: 78.8 OFF -> 113.6 t/s n=2 (1.44x) Qwen3.6-27B Q4_K_XL GPU: 78.8 OFF -> 113.6 t/s n=2 (1.44x) Qwen3.6-35B-A3B Q4 GPU: 192.3 OFF -> 223.2 t/s n=2 (1.16x) The 2B inflection is sharp. Skip auto-promote to draft-mtp when the identifier reports <2.0B params; users can still force via --spec-type or the Speculative Decoding toggle. Mirror the gate in the reload-skip check so a sub-2B reload-with-default does not bounce a spec-off backend. Chat-completions usage. llama-server's final SSE chunk emits both an OpenAI-style usage block and a custom timings block. timings.predicted_n is always populated, but usage.completion_tokens is zero on some server builds. The Studio chat UI computes generation t/s from meta.usage.completion_tokens / totalStreamTime, so a zero completion_tokens makes the UI fall back to wall-clock time (including SSE / proxy / template overhead) which dilutes MTP gains and makes ON look the same as OFF. Add _backfill_usage_from_timings: if usage.completion_tokens is missing or zero AND timings has predicted_n/prompt_n, synthesize a complete usage dict. Apply at the streaming metadata yield in generate_chat_completion and at the three accumulator/yield sites in generate_chat_completion_with_tools so per-iteration counts are not silently lost across tool calls. Tests cover both the gate (sub-2B skips, 2B+ promotes) and the backfill (zero usage filled, real usage preserved, empty timings passthrough). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: probe + emit legacy ngram-mod flags for pre-rename llama-server llama.cpp upstream renamed the ngram-mod tuning knobs: --draft-max -> --spec-ngram-mod-n-max (and --spec-draft-n-max) --draft-min -> --spec-ngram-mod-n-min (and --spec-draft-n-min) --spec-ngram-size-n -> --spec-ngram-mod-n-match The new names are real flags on post-rename builds and stub removal entries on the same builds (with description "argument has been removed"). Pre-rename builds only carry the legacy names as real flags. Studio was emitting the new names unconditionally, so a user running a pre-rename llama-server (e.g. an older prebuilt or a hand-installed binary) would see "unknown argument" errors when the ngram-mod path engages, or silent drop of the ngram knobs. Extend `probe_server_capabilities` to parse the help text into per-flag description blocks and tell real flags apart from removal stubs by the "argument has been removed" marker. Add three new probe fields: `ngram_mod_flavor` ("new" / "legacy" / None), `supports_ngram_mod`, and `spec_draft_n_max_flag` (the actual n_max flag the binary accepts). Cached by (path, mtime) the same way as `mtp_token`. Add `_build_ngram_mod_flags(caps, ...)` that picks the right flag set, returning [] when neither is usable so callers can drop ngram chaining entirely on minimal binaries. Wire both call sites to use the probe-driven flag set: - CPU/Mac MTP comma-chain (--spec-type ngram-mod,draft-mtp) emits legacy or new knobs as appropriate. If neither set is available, degrade to MTP-only (warn but still engage spec). - Standalone --spec-type ngram-mod branch uses the same helper. Tests cover post-rename detection, legacy detection, removal-stub discrimination, minimal-binary case, and all three branches of `_build_ngram_mod_flags` plus custom n_match/n_min/n_max values. Verified against three real binaries (Studio bundled 726704a, my build of 45b455e HEAD, and the MTP merge baseline 2555826) all correctly reporting ngram_mod_flavor=new. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: sub-3B MTP falls back to ngram-mod, not off Earlier sub-2B gate disabled speculative decoding entirely for tiny dense MTP models because the MTP draft head's per-token cost exceeds the acceptance savings at that scale. The "fully off" fallback was conservative -- ngram-mod has near-zero idle cost on diverse content and consistently outperforms both off and draft-mtp at sub-3B. Clean-methodology bench (each of 9 distinct prompts run once after two unrelated warmup prompts so the ngram-mod hash pool is realistically populated but never holds the exact deterministic output we're about to measure): Q4_K_XL on B200: 0.8B OFF=451 draft-mtp n=2=263 (0.58x) ngram-only=498 (1.10x) 2B OFF=377 draft-mtp n=2=308 (0.82x) ngram-only=369 (1.00x) 4B OFF=240 draft-mtp n=2=260 (1.08x) -- 4B+ wins with MTP Q4_K_XL on x86 48 cores: 0.8B OFF= 80 chained n=2= 69 (0.86x) ngram-only= 95 (1.19x) 2B OFF= 62 chained n=2= 51 (0.83x) ngram-only= 63 (1.01x) 4B OFF= 31 chained n=2= 41 (1.33x) Change: - Raise the MTP-skip threshold from 2.0B to 3.0B (2B falls below it). - When skipping the MTP head, fall back to --spec-type ngram-mod via the probe-driven _build_ngram_mod_flags helper. Works on both post-rename and pre-rename llama-server builds. - If the binary advertises neither ngram-mod flavor, fall back to spec-off (older binaries that don't support ngram-mod at all). - Mirror the same fallback in _already_in_target_state so a sub-3B reload-with-default does not bounce a ngram-mod backend. Tests updated: monkeypatch probe_server_capabilities so the gate behavior is deterministic regardless of which llama-server happens to be on the host. +1 new test for the "binary has no ngram-mod support" branch; renamed prior 2B/0.8B tests to reflect new semantics. This generalizes the size gate to be probe-driven instead of a hard "disable spec" branch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: 5-mode Speculative Decoding dropdown (Auto / MTP / Ngram / MTP+Ngram / Off) Replace the Chat Settings Speculative Decoding on/off Switch with a 5-option Select. Auto preserves today's platform-aware resolver (MTP on MTP GGUFs, ngram-mod fallback for sub-3B, --spec-default for non-MTP). The other 3 modes force the user's choice on BOTH GPU and CPU: MTP emits draft-mtp only (no ngram chain on CPU), Ngram emits ngram-mod only, MTP+Ngram emits the ngram-mod,draft-mtp chain on both platforms. Off is the existing fully-off state, kept so the Switch's "disable" capability isn't lost. Backend - New module-level _canonicalize_spec_mode(value) maps any accepted input (canonical, legacy "default" / "draft-mtp" / "ngram-mod" / "ngram-simple", or comma-chained "ngram-mod,draft-mtp") onto one of auto / mtp / ngram / mtp+ngram / off / ngram-simple / None. Lets external callers and old persisted UI state round-trip without breaking. - LlamaCppBackend grows a _requested_spec_mode field + requested_spec_mode property storing the canonical UI mode the user requested. Status responses round-trip this instead of the resolved internal flag, so the dropdown restores the picked value after reload / refresh (Auto on a 27B MTP GGUF resolves to draft-mtp internally but the dropdown stays on "Auto"). - The resolver block in load_model is extracted into a unit-testable _build_speculative_flags method. Forced MTP / MTP+Ngram on a sub-3B or non-MTP GGUF logs a warning and engages anyway (user override > the Auto-path sub-3B fallback). - _already_in_target_state and routes/inference._request_matches_loaded_settings now compare canonical-requested mode, dropping the old auto-promotion mirror. spec_draft_n_max still gates on the resolved spec so Auto + a changed n_max still bounces a reload. Frontend - chat-settings-sheet.tsx: Switch swapped for Select modeled on the KV Cache Dtype Select. Items: Auto / MTP / Ngram / MTP+Ngram / Off. Draft Tokens input only visible when speculativeType is "mtp" or "mtp+ngram". - chat-runtime-store.ts: initial value flips from "default" to "auto". - use-chat-model-runtime.ts normalizeSpeculativeType mirrors the backend canonicaliser so persisted "default" / "draft-mtp" / "ngram-mod" / chain values hydrate to the right dropdown option. - types/api.ts: docs the canonical wire vocabulary. Tests - 53 new assertions in test_llama_cpp_mtp_detection.py: full _canonicalize_spec_mode table, a 23-row resolver matrix across (requested mode) x (GPU/CPU) x (model size class), plus n_max override, user-extra-args precedence, requested-mode round-trip, and graceful degrade on an outdated llama-server without an MTP token. - 165 existing backend tests still green. 218 total in the MTP / server-args / reload-inheritance suite. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: reset Speculative Decoding to Auto on model switch When the user switches from model A to a different model B, clear the runtime store's speculativeType + specDraftNMax (and their loaded* shadows). The new load request then carries null, the backend canonicalises that to "auto", and its platform-aware resolver runs fresh for the new model. Without this, a non-MTP model loaded with "Off" carried the Off choice into a subsequent MTP load, suppressing MTP auto-promotion (and the sub-3B ngram-mod fallback) until the user manually opened settings and flipped the dropdown back to Auto. The clean-sweep deep probe caught it as anomaly A-1. The reset only fires when currentCheckpoint != modelId, so a same-model reapply or forceReload still honours the user's current spec choice. End-to-end probe on Qwen3.5-4B-GGUF (non-MTP, Off) -> Qwen3.5-0.8B-MTP confirms: dropdown shows Auto, /api/inference/status returns speculative_type=auto, studio.log shows the Auto sub-3B fallback emitted --spec-type ngram-mod. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 610 +++++++++++++--- .../core/inference/llama_server_args.py | 2 + studio/backend/models/inference.py | 51 +- studio/backend/routes/inference.py | 31 +- .../tests/test_gguf_reload_inheritance.py | 1 + .../tests/test_llama_cpp_mtp_detection.py | 661 +++++++++++++++++- .../src/features/chat/chat-settings-sheet.tsx | 91 ++- .../chat/hooks/use-chat-model-runtime.ts | 52 +- .../chat/stores/chat-runtime-store.ts | 13 +- .../frontend/src/features/chat/types/api.ts | 18 + 10 files changed, 1385 insertions(+), 145 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 1791f06676..134b6574b0 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -479,6 +479,122 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: return False +def _build_ngram_mod_flags( + caps: Optional[dict], + n_match: int = 24, + n_min: int = 48, + n_max: int = 64, +) -> list[str]: + """Emit the right ngram-mod knob flags for the running llama-server. + + Post-rename builds expose ``--spec-ngram-mod-n-{match,min,max}``; + pre-rename builds expose the legacy ``--spec-ngram-size-n`` / + ``--draft-min`` / ``--draft-max``. ``caps`` comes from + ``probe_server_capabilities``; ``ngram_mod_flavor`` tells us which + set is real (vs a removal-stub entry). Returns ``[]`` when neither + set is available so the caller can drop ngram-mod entirely. + """ + flavor = caps.get("ngram_mod_flavor") if caps else None + if flavor == "new": + return [ + "--spec-ngram-mod-n-match", + str(n_match), + "--spec-ngram-mod-n-min", + str(n_min), + "--spec-ngram-mod-n-max", + str(n_max), + ] + if flavor == "legacy": + # Legacy llama.cpp before the spec arg rename: same knobs lived + # under --spec-ngram-size-n (lookup length) and the generic + # --draft-min / --draft-max (ngram size N range). + return [ + "--spec-ngram-size-n", + str(n_match), + "--draft-min", + str(n_min), + "--draft-max", + str(n_max), + ] + return [] + + +# Canonical Speculative Decoding modes exposed by the Studio chat UI. +# The dropdown renders five options (auto, mtp, ngram, mtp+ngram, off); +# the load API also accepts legacy values that the original Switch and +# external callers emit (default, draft-mtp, ngram-mod, ngram-simple). +_CANONICAL_SPEC_MODES = {"auto", "mtp", "ngram", "mtp+ngram", "off", "ngram-simple"} +_LEGACY_SPEC_MODE_MAP = { + "default": "auto", + "draft-mtp": "mtp", + "ngram-mod": "ngram", +} + + +def _canonicalize_spec_mode(value): + """Map any accepted ``speculative_type`` input onto a canonical mode. + + Returns one of ``auto``, ``mtp``, ``ngram``, ``mtp+ngram``, ``off``, + ``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``). + Unknown strings collapse to ``auto`` so a stale UI value or typo + falls back to the safe platform-aware path. + """ + if value is None: + return None + if not isinstance(value, str): + return None + stripped = value.strip().lower() + if not stripped: + return None + if stripped in _CANONICAL_SPEC_MODES: + return stripped + if stripped in _LEGACY_SPEC_MODE_MAP: + return _LEGACY_SPEC_MODE_MAP[stripped] + # llama.cpp comma-chains are emitted by old persisted state e.g. + # "ngram-mod,draft-mtp"; collapse the most common one explicitly. + pieces = [p.strip() for p in stripped.split(",") if p.strip()] + has_mtp = any(p in ("mtp", "draft-mtp") for p in pieces) + has_ngram = any(p in ("ngram", "ngram-mod") for p in pieces) + if has_mtp and has_ngram: + return "mtp+ngram" + if has_mtp: + return "mtp" + if has_ngram: + return "ngram" + return "auto" + + +def _backfill_usage_from_timings(usage, timings): + """Synthesize ``usage`` from llama-server's ``timings`` when the + OpenAI-style usage block is missing or reports zero tokens. + + The Studio chat UI computes generation t/s from + ``meta.usage.completion_tokens / totalStreamTime``. llama-server + always populates ``timings.predicted_n`` (true decoded count) and + ``timings.prompt_n``, but the ``usage`` field on the final SSE chunk + can be absent or zero on some server builds / streaming + configurations, which makes the UI fall back to wall-clock t/s and + dilute speculative-decoding speedups. + """ + if not timings: + return usage + if usage and usage.get("completion_tokens"): + return usage + predicted_n = timings.get("predicted_n") + prompt_n = timings.get("prompt_n") + if predicted_n is None and prompt_n is None: + return usage + out = dict(usage or {}) + if not out.get("completion_tokens") and predicted_n is not None: + out["completion_tokens"] = predicted_n + if not out.get("prompt_tokens") and prompt_n is not None: + out["prompt_tokens"] = prompt_n + out["total_tokens"] = int(out.get("prompt_tokens") or 0) + int( + out.get("completion_tokens") or 0 + ) + return out + + class LlamaCppBackend: """ Manages a llama-server subprocess for GGUF model inference. @@ -513,6 +629,15 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None self._reasoning_default: bool = True self._speculative_type: Optional[str] = None + # Canonical UI-facing mode the user requested: one of + # ``auto``/``mtp``/``ngram``/``mtp+ngram``/``off``/``ngram-simple``. + # Round-tripped through the status API so the dropdown reflects + # the picked mode rather than the resolved internal flag set + # (auto on a 27B MTP GGUF resolves to draft-mtp but the dropdown + # should still read "Auto"). + self._requested_spec_mode: Optional[str] = None + # User-supplied --spec-draft-n-max override (None = platform default). + self._spec_draft_n_max: Optional[int] = None # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None self._n_kv_heads: Optional[int] = None @@ -793,6 +918,17 @@ class LlamaCppBackend: def speculative_type(self) -> Optional[str]: return self._speculative_type + @property + def requested_spec_mode(self) -> Optional[str]: + """Canonical UI-facing mode the user requested (see field doc).""" + return self._requested_spec_mode + + @property + def spec_draft_n_max(self) -> Optional[int]: + """User --spec-draft-n-max override active on the load, or None + when the platform default (6 GPU / 3 CPU) is in effect.""" + return self._spec_draft_n_max + # ── Binary discovery ────────────────────────────────────────── @staticmethod @@ -919,11 +1055,31 @@ class LlamaCppBackend: cls, binary: Optional[str] = None ) -> dict[str, object]: """Parse `llama-server --help` for feature flags. Returns - {found, mtp_token, supports_mtp}. mtp_token is "draft-mtp" - (older) or "mtp" (renamed upstream), or None.""" + {found, mtp_token, supports_mtp, ngram_mod_flavor, + supports_ngram_mod, spec_draft_n_max_flag}. + + ``ngram_mod_flavor`` is ``"new"`` when the binary exposes the + post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` as + real args, ``"legacy"`` when only the pre-rename + ``--spec-ngram-size-n / --draft-min / --draft-max`` are real + (the rename ships with stub removal entries for the legacy + names; we tell stubs apart by the "argument has been removed" + description), or ``None`` if neither set is usable. + + ``spec_draft_n_max_flag`` is the actual flag name the binary + accepts: ``--spec-draft-n-max`` on post-rename builds, or + ``--draft-max`` on legacy. ``None`` means n_max cannot be set. + """ bin_path = binary or cls._find_llama_server_binary() if not bin_path or not Path(bin_path).is_file(): - return {"found": False, "mtp_token": None, "supports_mtp": False} + return { + "found": False, + "mtp_token": None, + "supports_mtp": False, + "ngram_mod_flavor": None, + "supports_ngram_mod": False, + "spec_draft_n_max_flag": None, + } try: mtime = int(Path(bin_path).stat().st_mtime) except OSError: @@ -934,6 +1090,8 @@ class LlamaCppBackend: return cached mtp_token: Optional[str] = None + ngram_mod_flavor: Optional[str] = None + spec_draft_n_max_flag: Optional[str] = None try: result = subprocess.run( [bin_path, "--help"], @@ -943,6 +1101,52 @@ class LlamaCppBackend: check = False, ) help_text = (result.stdout or "") + "\n" + (result.stderr or "") + # Split into per-flag blocks: each --flag line plus its + # indented continuation lines, so the "argument has been + # removed" description sits with its flag. + blocks: dict[str, str] = {} + current_flags: list[str] = [] + current_desc: list[str] = [] + for line in help_text.splitlines(): + stripped = line.strip() + if stripped.startswith("-") and not line.startswith(" "): + # New flag line; flush previous. + if current_flags: + desc = " ".join(current_desc) + for f in current_flags: + blocks[f] = desc + current_flags = [] + current_desc = [stripped] + # Extract long-form flag tokens from the DECLARATION + # prefix only (comma-separated aliases). Stop at the + # first token that isn't itself a flag, so flag + # references inside descriptions are ignored. + for tok in re.split(r"[,\s]+", stripped): + if tok.startswith("--") and re.match( + r"--[A-Za-z][A-Za-z0-9_-]*$", tok + ): + current_flags.append(tok) + elif tok.startswith("-") and len(tok) > 1: + # short alias like -fa; keep scanning aliases. + continue + else: + # First non-flag token marks end of decl. + break + else: + current_desc.append(stripped) + if current_flags: + desc = " ".join(current_desc) + for f in current_flags: + blocks[f] = desc + + def _is_real(flag: str) -> bool: + """True if the flag exists AND is not a removal stub.""" + desc = blocks.get(flag) + if desc is None: + return False + return "argument has been removed" not in desc + + # MTP token detection from --spec-type line. spec_line = "" for line in help_text.splitlines(): if "--spec-type" in line: @@ -953,6 +1157,30 @@ class LlamaCppBackend: mtp_token = "draft-mtp" elif re.search(r"[|,\[]mtp[|,\]]", spec_line): mtp_token = "mtp" + + # ngram-mod flag flavor. Post-rename builds advertise both + # the new args (real) and the legacy ones (stubs); pre-rename + # builds only have the legacy ones as real. + new_ngram_real = ( + _is_real("--spec-ngram-mod-n-match") + and _is_real("--spec-ngram-mod-n-min") + and _is_real("--spec-ngram-mod-n-max") + ) + legacy_ngram_real = ( + _is_real("--spec-ngram-size-n") + and _is_real("--draft-max") + and _is_real("--draft-min") + ) + if new_ngram_real: + ngram_mod_flavor = "new" + elif legacy_ngram_real: + ngram_mod_flavor = "legacy" + + # n_max flag: prefer post-rename, fall back to legacy. + if _is_real("--spec-draft-n-max"): + spec_draft_n_max_flag = "--spec-draft-n-max" + elif _is_real("--draft-max"): + spec_draft_n_max_flag = "--draft-max" except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") @@ -960,6 +1188,9 @@ class LlamaCppBackend: "found": True, "mtp_token": mtp_token, "supports_mtp": mtp_token is not None, + "ngram_mod_flavor": ngram_mod_flavor, + "supports_ngram_mod": ngram_mod_flavor is not None, + "spec_draft_n_max_flag": spec_draft_n_max_flag, } cls._capability_cache[cache_key] = info return info @@ -2258,6 +2489,7 @@ class LlamaCppBackend: chat_template_override: Optional[str] = None, cache_type_kv: Optional[str] = None, speculative_type: Optional[str] = None, + spec_draft_n_max: Optional[int] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused n_parallel: int = 1, @@ -2289,6 +2521,7 @@ class LlamaCppBackend: n_ctx = n_ctx, cache_type_kv = cache_type_kv, speculative_type = speculative_type, + spec_draft_n_max = spec_draft_n_max, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -2636,98 +2869,16 @@ class LlamaCppBackend: # (llama.cpp #22673). Auto-enabled via nextn_predict_layers, # fallback to -MTP in name. GPU: MTP-only. CPU/Mac: chain # with ngram-mod. See unsloth.ai/docs/models/qwen3.6#mtp-guide. - _valid_spec_types = {"ngram-simple", "ngram-mod", "draft-mtp"} - normalized_spec = ( - speculative_type.lower().strip() if speculative_type else None + spec_flags = self._build_speculative_flags( + speculative_type = speculative_type, + spec_draft_n_max = spec_draft_n_max, + extra_args = extra_args, + model_identifier = model_identifier, + model_path = model_path, + gpus = bool(gpus), + binary = binary, ) - is_mtp_model = bool(self._nextn_predict_layers) or ( - _is_mtp_model_name(model_identifier, model_path) - ) - user_owns_spec_type = _extra_args_set_spec_type(extra_args) - # Auto-promote unset/"default" to draft-mtp on MTP GGUFs. - # llama.cpp #22673: MTP is compatible with mmproj, so the - # vision gate previously here was wrong. - if ( - is_mtp_model - and not user_owns_spec_type - and normalized_spec in (None, "", "default") - ): - normalized_spec = "draft-mtp" - if user_owns_spec_type: - # User --spec-type wins (it accumulates if repeated). - normalized_spec = None - self._speculative_type = None - if normalized_spec and normalized_spec != "off": - if normalized_spec == "default": - cmd.append("--spec-default") - self._speculative_type = "default" - elif normalized_spec == "draft-mtp": - # Probe binary; fail gracefully on outdated prebuilts. - # Use whichever token the binary advertises - # (older: draft-mtp; renamed upstream: mtp). - caps = self.probe_server_capabilities(binary) - mtp_token = caps.get("mtp_token") if caps else None - if not mtp_token: - logger.warning( - "MTP GGUF detected but llama-server lacks " - "--spec-type mtp/draft-mtp; run " - "`unsloth studio update`. Loading without " - "speculative decoding." - ) - self._speculative_type = None - else: - if gpus: - cmd.extend( - [ - "--spec-type", - mtp_token, - "--spec-draft-n-max", - "6", - ] - ) - else: - # CPU/Mac: chain ngram-mod + MTP in one - # comma-separated --spec-type (not repeated). - # ngram-mod knobs match llama.cpp defaults - # (n-match 24, n-min 48, n-max 64). - cmd.extend( - [ - "--spec-type", - f"ngram-mod,{mtp_token}", - "--spec-draft-n-max", - "3", - "--spec-ngram-mod-n-match", - "24", - "--spec-ngram-mod-n-min", - "48", - "--spec-ngram-mod-n-max", - "64", - ] - ) - self._speculative_type = "draft-mtp" - logger.info( - f"Spec decoding: {mtp_token} ({'GPU' if gpus else 'CPU/Mac'})" - ) - elif normalized_spec in _valid_spec_types: - cmd.extend(["--spec-type", normalized_spec]) - if normalized_spec == "ngram-mod": - # llama.cpp defaults; legacy --spec-ngram-size-n - # / --draft-{min,max} were removed for ngram-mod. - cmd.extend( - [ - "--spec-ngram-mod-n-match", - "24", - "--spec-ngram-mod-n-min", - "48", - "--spec-ngram-mod-n-max", - "64", - ] - ) - self._speculative_type = normalized_spec - else: - self._speculative_type = None - else: - self._speculative_type = None + cmd.extend(spec_flags) # Apply custom chat template override if provided self._chat_template_override = chat_template_override @@ -3059,6 +3210,220 @@ class LlamaCppBackend: ) return True + def _build_speculative_flags( + self, + *, + speculative_type: Optional[str], + spec_draft_n_max: Optional[int], + extra_args: Optional[List[str]], + model_identifier: str, + model_path: Optional[str], + gpus: bool, + binary: Optional[str], + ) -> List[str]: + """Return the llama-server flag list for the requested spec mode. + + Side effects: sets ``self._speculative_type`` (resolved internal + emit), ``self._requested_spec_mode`` (canonical UI mode for the + status round-trip), and ``self._spec_draft_n_max`` (user override + only; None when the platform default applies). + + Speculative decoding (n-gram self-speculation, zero VRAM cost): + ngram-mod uses a ~16 MB shared hash pool, constant memory / + complexity, variable draft lengths. Helps most when the model + repeats existing text (code refactor, summarisation, reasoning). + For general chat with low repetition, overhead is ~5 ms. + + Benchmarks from upstream llama.cpp speculative-decoding PRs: + Scenario | Without | With | Speedup + gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x + Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x + gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x + + Sub-3B dense MTP regresses vs spec-off because the draft head's + per-token cost exceeds the acceptance savings at this scale. + Q4_K_XL clean bench (each prompt once after an unrelated warmup) + on B200 + x86 CPU: + 0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x + 2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x + 0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x + 2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x + 4B+ GPU/CPU: spec on is a net win (1.08x-1.46x). + Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on + diverse content); forced MTP variants engage anyway and just log + a warning per the user's choice. + """ + flags: List[str] = [] + # Reset; emit branches re-set on the resolved emission. + self._spec_draft_n_max = None + self._speculative_type = None + + # Canonical UI-facing requested mode: auto / mtp / ngram / + # mtp+ngram / off / ngram-simple. Legacy values are mapped via + # _canonicalize_spec_mode (default->auto, draft-mtp->mtp, + # ngram-mod->ngram, "ngram-mod,draft-mtp"->mtp+ngram). + canonical_mode = _canonicalize_spec_mode(speculative_type) + is_mtp_model = bool(self._nextn_predict_layers) or ( + _is_mtp_model_name(model_identifier, model_path) + ) + user_owns_spec_type = _extra_args_set_spec_type(extra_args) + _mtp_size_b = _extract_model_size_b(model_identifier) + _mtp_too_small = _mtp_size_b is not None and _mtp_size_b < 3.0 + + if user_owns_spec_type: + # User --spec-type in extra_args wins outright; suppress + # auto-emit so we don't emit a duplicate / conflicting + # spec block. Record requested mode as None. + self._requested_spec_mode = None + return flags + + effective_mode = canonical_mode or "auto" + self._requested_spec_mode = effective_mode + + def _resolved_draft_n_max() -> int: + # User override wins; else platform default (the B200 / x86 + # clean-sweep sweet spot from PR #5582 is n=2 GPU, n=3 CPU; + # raising past 3 starts to regress on essay-style + # low-acceptance prompts). + if spec_draft_n_max is not None: + n = int(spec_draft_n_max) + self._spec_draft_n_max = n + return n + return 2 if gpus else 3 + + def _emit_mtp(*, chain_ngram: bool) -> bool: + """Append --spec-type mtp[/draft-mtp][,ngram-mod] + n-max.""" + caps = self.probe_server_capabilities(binary) + mtp_token = caps.get("mtp_token") if caps else None + if not mtp_token: + logger.warning( + "Requested MTP speculative decoding but " + "llama-server lacks --spec-type mtp/draft-mtp; " + "run `unsloth studio update`. Loading without " + "speculative decoding." + ) + return False + draft_n_max = _resolved_draft_n_max() + n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max" + if chain_ngram: + ngram_knobs = _build_ngram_mod_flags(caps) + if ngram_knobs: + spec_value = f"ngram-mod,{mtp_token}" + else: + logger.warning( + "llama-server lacks ngram-mod tuning " + "flags; loading MTP only (no ngram chain)" + ) + spec_value = mtp_token + flags.extend( + [ + "--spec-type", + spec_value, + n_max_flag, + str(draft_n_max), + ] + ) + flags.extend(ngram_knobs) + else: + flags.extend( + [ + "--spec-type", + mtp_token, + n_max_flag, + str(draft_n_max), + ] + ) + self._speculative_type = "draft-mtp" + chain_label = "chained ngram-mod" if chain_ngram else "MTP-only" + logger.info(f"Spec decoding: {mtp_token} ({chain_label})") + return True + + def _emit_ngram_mod() -> bool: + """Append --spec-type ngram-mod + flag-set knobs.""" + ngram_caps = self.probe_server_capabilities(binary) + ngram_knobs = _build_ngram_mod_flags(ngram_caps) + flags.extend(["--spec-type", "ngram-mod"]) + if not ngram_knobs: + logger.warning( + "llama-server lacks ngram-mod tuning " + "flags; loading without --spec-ngram-mod-* knobs" + ) + flags.extend(ngram_knobs) + self._speculative_type = "ngram-mod" + logger.info("Spec decoding: ngram-mod") + return True + + if effective_mode == "off": + return flags # nothing to emit + if effective_mode == "ngram-simple": + flags.extend(["--spec-type", "ngram-simple"]) + self._speculative_type = "ngram-simple" + return flags + if effective_mode == "ngram": + _emit_ngram_mod() + return flags + if effective_mode == "mtp": + if _mtp_too_small: + logger.warning( + f"Forcing MTP on a {_mtp_size_b:.1f}B model; " + "the bench shows draft-mtp regresses below 3B. " + "Engaging anyway (user override)." + ) + elif not is_mtp_model: + logger.warning( + "Forcing MTP on a non-MTP GGUF; llama-server may " + "fall back to spec-off if no nextn head is present. " + "Engaging anyway (user override)." + ) + _emit_mtp(chain_ngram = False) + return flags + if effective_mode == "mtp+ngram": + if _mtp_too_small: + logger.warning( + f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; " + "the bench shows the chain regresses below 3B. " + "Engaging anyway (user override)." + ) + elif not is_mtp_model: + logger.warning( + "Forcing MTP+Ngram on a non-MTP GGUF; llama-server " + "may fall back to ngram-only if no nextn head is " + "present. Engaging anyway (user override)." + ) + _emit_mtp(chain_ngram = True) + return flags + + # effective_mode == "auto": today's promotion path. llama.cpp + # #22673: MTP is compatible with mmproj, so there's no vision gate. + if is_mtp_model and not _mtp_too_small: + # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. + _emit_mtp(chain_ngram = not gpus) + elif is_mtp_model and _mtp_too_small: + # Sub-3B fallback: drop the MTP draft head, keep ngram-mod + # when the binary supports it. + _small_caps = self.probe_server_capabilities(binary) + if _small_caps.get("supports_ngram_mod"): + logger.info( + f"MTP GGUF detected but model size {_mtp_size_b:.1f}B " + "is below the 3B speedup threshold; using ngram-mod " + "only (zero-VRAM, no draft head). Override via " + "--spec-type or the Studio Speculative Decoding " + "dropdown." + ) + _emit_ngram_mod() + else: + logger.info( + f"MTP GGUF detected but model size {_mtp_size_b:.1f}B " + "is below the 3B speedup threshold and the bundled " + "llama-server does not advertise ngram-mod; " + "auto-disabling speculative decoding." + ) + else: + # Non-MTP model: let llama-server choose its default strategy. + flags.append("--spec-default") + self._speculative_type = "default" + return flags + def _already_in_target_state( self, *, @@ -3071,6 +3436,7 @@ class LlamaCppBackend: extra_args: Optional[List[str]], is_vision: bool, gguf_path: Optional[str] = None, + spec_draft_n_max: Optional[int] = None, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -3107,18 +3473,28 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # Mirror load_model's auto-promotion. Vision is no longer a - # spec blocker (llama.cpp #22673: MTP is compatible with mmproj). - raw_spec = _norm(speculative_type) - req_spec = raw_spec or "off" + # Compare on the canonical UI-facing mode the user requested. + # When extra_args carries --spec-type, the route-layer code paths + # bypass the dropdown anyway and the backend stores + # _requested_spec_mode = None; the request mirrors that by + # canonicalising to None. + if _extra_args_set_spec_type(extra_args): + req_mode = None + else: + req_mode = _canonicalize_spec_mode(speculative_type) or "auto" + backend_mode = self._requested_spec_mode + if req_mode != backend_mode: + return False + + # spec_draft_n_max only matters when an MTP variant is actually + # engaged. Compare on the resolved spec rather than the requested + # mode so an Auto request that auto-promoted to draft-mtp under + # the hood still bounces a reload when the user changes n_max. if ( - raw_spec in (None, "default") - and _is_mtp_model_name(model_identifier, gguf_path) - and not _extra_args_set_spec_type(extra_args) + self._speculative_type == "draft-mtp" + and spec_draft_n_max is not None + and int(spec_draft_n_max) != (self._spec_draft_n_max or 0) ): - req_spec = "draft-mtp" - backend_spec = _norm(self._speculative_type) or "off" - if req_spec != backend_spec: return False if (self._chat_template_override or None) != (chat_template_override or None): @@ -3187,6 +3563,8 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._speculative_type = None + self._requested_spec_mode = None + self._spec_draft_n_max = None self._n_layers = None self._n_kv_heads = None self._n_kv_heads_by_layer = None @@ -3816,6 +4194,9 @@ class LlamaCppBackend: if _stream_done: break # exit outer for if _metadata_usage or _metadata_timings: + _metadata_usage = _backfill_usage_from_timings( + _metadata_usage, _metadata_timings + ) yield { "type": "metadata", "usage": _metadata_usage, @@ -4265,7 +4646,10 @@ class LlamaCppBackend: } ) # Accumulate tokens and timing from this iteration - _fu_r = _iter_usage or {} + _fu_r = ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) + or {} + ) _accumulated_completion_tokens += _fu_r.get( "completion_tokens", 0 ) @@ -4277,7 +4661,10 @@ class LlamaCppBackend: # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} - _fu = _iter_usage or {} + _fu = ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) + or {} + ) _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens @@ -4367,7 +4754,10 @@ class LlamaCppBackend: ) if content_accum: yield {"type": "content", "text": content_accum} - _fu = _iter_usage or {} + _fu = ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) + or {} + ) _fc = _fu.get("completion_tokens", 0) _fp = _fu.get("prompt_tokens", 0) _tc = _fc + _accumulated_completion_tokens @@ -4401,9 +4791,9 @@ class LlamaCppBackend: return # ── Execute tool calls ── - _accumulated_completion_tokens += (_iter_usage or {}).get( - "completion_tokens", 0 - ) + _accumulated_completion_tokens += ( + _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} + ).get("completion_tokens", 0) _it = _iter_timings or {} _accumulated_predicted_ms += _it.get("predicted_ms", 0) _accumulated_predicted_n += _it.get("predicted_n", 0) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 572ac2ceda..d8b7eb383e 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -148,6 +148,8 @@ _SPEC_FLAGS: frozenset[str] = frozenset( # MTP path (llama.cpp #22673). "--spec-draft-n-max", "--spec-draft-n-min", + "--spec-draft-p-min", + "--spec-draft-p-split", "--spec-ngram-mod-n-match", "--spec-ngram-mod-n-min", "--spec-ngram-mod-n-max", diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 99d1df37b6..e32d134628 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -70,7 +70,28 @@ class LoadRequest(BaseModel): ) speculative_type: Optional[str] = Field( None, - description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.", + description = ( + "Speculative decoding mode for GGUF models. Canonical values: " + "'auto' (platform-aware: MTP on MTP GGUFs, ngram-mod fallback " + "for sub-3B), 'mtp' (force draft-mtp only on both GPU and CPU), " + "'ngram' (force ngram-mod only), 'mtp+ngram' (force " + "ngram-mod+draft-mtp chain on both platforms), 'off' (disabled). " + "Legacy values 'default' (-> auto), 'draft-mtp' (-> mtp), " + "'ngram-mod' (-> ngram), and 'ngram-simple' (kept as-is) are " + "still accepted. Ignored for non-GGUF and vision models." + ), + ) + spec_draft_n_max: Optional[int] = Field( + None, + ge = 1, + le = 16, + description = ( + "Max draft tokens per step for MTP speculative decoding " + "(--spec-draft-n-max). Defaults to 2 on GPU and 3 on CPU/Mac " + "when unset (upstream-bench sweet spot for dense Qwen3.6 MTP " + "quants). Only applied when speculative_type resolves to " + "'mtp' or 'mtp+ngram'." + ), ) llama_extra_args: Optional[List[str]] = Field( None, @@ -218,7 +239,19 @@ class LoadResponse(BaseModel): ) speculative_type: Optional[str] = Field( None, - description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled", + description = ( + "Canonical UI-facing requested speculative decoding mode " + "('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / " + "'ngram-simple'), round-tripped from the original LoadRequest " + "via _canonicalize_spec_mode. None when no model is loaded." + ), + ) + spec_draft_n_max: Optional[int] = Field( + None, + description = ( + "Active --spec-draft-n-max for MTP speculative decoding, or " + "None when the platform default is in effect." + ), ) @@ -340,7 +373,19 @@ class InferenceStatusResponse(BaseModel): ) speculative_type: Optional[str] = Field( None, - description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled", + description = ( + "Canonical UI-facing requested speculative decoding mode " + "('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / " + "'ngram-simple'), round-tripped from the original LoadRequest. " + "None when no model is loaded." + ), + ) + spec_draft_n_max: Optional[int] = Field( + None, + description = ( + "Active --spec-draft-n-max for MTP speculative decoding, or " + "None when the platform default is in effect." + ), ) llama_cpp_supports_mtp: bool = Field( True, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2ed3315f56..2f095cb823 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -117,6 +117,7 @@ try: LlamaCppBackend, _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_T_MAX_PREDICT_MS, + _canonicalize_spec_mode, _hf_offline_if_dns_dead, detect_reasoning_flags, ) @@ -143,6 +144,7 @@ except ImportError: LlamaCppBackend, _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_T_MAX_PREDICT_MS, + _canonicalize_spec_mode, _hf_offline_if_dns_dead, detect_reasoning_flags, ) @@ -441,12 +443,17 @@ def _request_matches_loaded_settings( # spec on ``not is_vision``), so treat the request as ``off`` against # the backend's ``None`` to avoid forcing a redundant reload. if llama_backend.is_vision: - req_spec = "off" + req_mode = "off" else: - req_spec = _normalise_settings_str(request.speculative_type) or "off" - backend_spec = _normalise_settings_str(llama_backend.speculative_type) or "off" - if req_spec != backend_spec: + req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto" + backend_mode = llama_backend.requested_spec_mode or "auto" + if req_mode != backend_mode: return False + # spec_draft_n_max only matters when an MTP variant is engaged; None + # means "platform default" and matches whatever the backend chose. + if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None: + if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0): + return False if (request.chat_template_override or None) != ( llama_backend.chat_template_override or None ): @@ -584,7 +591,8 @@ async def load_model( reasoning_always_on = llama_backend.reasoning_always_on, supports_preserve_thinking = llama_backend.supports_preserve_thinking, chat_template = llama_backend.chat_template, - speculative_type = llama_backend.speculative_type, + speculative_type = llama_backend.requested_spec_mode, + spec_draft_n_max = llama_backend.spec_draft_n_max, ) else: if ( @@ -724,7 +732,10 @@ async def load_model( llama_backend.extra_args, strip_context = "max_seq_length" in fields_set, strip_cache = "cache_type_kv" in fields_set, - strip_spec = "speculative_type" in fields_set, + strip_spec = ( + "speculative_type" in fields_set + or "spec_draft_n_max" in fields_set + ), strip_template = "chat_template_override" in fields_set, ) try: @@ -765,6 +776,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + spec_draft_n_max = request.spec_draft_n_max, n_parallel = _n_parallel, extra_args = extra_llama_args, ) @@ -788,6 +800,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + spec_draft_n_max = request.spec_draft_n_max, n_parallel = _n_parallel, extra_args = extra_llama_args, ) @@ -846,7 +859,8 @@ async def load_model( supports_tools = llama_backend.supports_tools, cache_type_kv = llama_backend.cache_type_kv, chat_template = llama_backend.chat_template, - speculative_type = llama_backend.speculative_type, + speculative_type = llama_backend.requested_spec_mode, + spec_draft_n_max = llama_backend.spec_draft_n_max, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -1345,7 +1359,8 @@ async def get_status( native_context_length = llama_backend.native_context_length, cache_type_kv = llama_backend.cache_type_kv, chat_template_override = llama_backend.chat_template_override, - speculative_type = llama_backend.speculative_type, + speculative_type = llama_backend.requested_spec_mode, + spec_draft_n_max = llama_backend.spec_draft_n_max, llama_cpp_supports_mtp = _supports_mtp, llama_cpp_prebuilt_stale = _stale, llama_cpp_installed_tag = _installed_tag, diff --git a/studio/backend/tests/test_gguf_reload_inheritance.py b/studio/backend/tests/test_gguf_reload_inheritance.py index 4b0b450cb0..1a663725d9 100644 --- a/studio/backend/tests/test_gguf_reload_inheritance.py +++ b/studio/backend/tests/test_gguf_reload_inheritance.py @@ -78,6 +78,7 @@ def _loaded_backend(**overrides): backend._requested_n_ctx = 8192 backend._cache_type_kv = None backend._speculative_type = None + backend._requested_spec_mode = "auto" backend._chat_template_override = None backend._is_vision = False backend._extra_args = None diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 7da633201f..4a8276adc0 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -52,6 +52,9 @@ import pytest from core.inference.llama_cpp import ( LlamaCppBackend, + _backfill_usage_from_timings, + _build_ngram_mod_flags, + _canonicalize_spec_mode, _extra_args_set_spec_type, _is_mtp_model_name, ) @@ -186,6 +189,10 @@ def _mtp_backend(**overrides): backend._requested_n_ctx = 8192 backend._cache_type_kv = None backend._speculative_type = "draft-mtp" + # Default fixture simulates Auto having auto-promoted to draft-mtp. + # Individual tests override _requested_spec_mode when they want a + # forced mode or the user---spec-type-extra-args path. + backend._requested_spec_mode = "auto" backend._chat_template_override = None backend._is_vision = False backend._extra_args = None @@ -233,9 +240,16 @@ def test_already_in_target_state_matches_when_request_uses_default_for_mtp_model ) -def test_already_in_target_state_non_mtp_model_unaffected(): - # Promotion is gated on the name; non-MTP must still mismatch req=None. - backend = _mtp_backend(_model_identifier = "unsloth/Qwen3.6-27B-GGUF") +def test_already_in_target_state_auto_request_matches_auto_backend_for_non_mtp_model(): + # Under the requested-mode round-trip model, Auto requested against an + # Auto-recorded backend matches regardless of model name. The underlying + # resolved emission (--spec-default vs draft-mtp) is handled by the + # backend's own load path and reflected in _speculative_type; the + # short-circuit comparison only cares whether the *intent* changed. + backend = _mtp_backend( + _model_identifier = "unsloth/Qwen3.6-27B-GGUF", + _speculative_type = "default", + ) assert ( backend._already_in_target_state( gguf_path = None, @@ -248,7 +262,7 @@ def test_already_in_target_state_non_mtp_model_unaffected(): extra_args = None, is_vision = False, ) - is False + is True ) @@ -308,6 +322,7 @@ def test_already_in_target_state_user_spec_type_override_matches_clean_backend() # User --spec-type none suppressed auto-MTP; repeat /load must not re-promote. backend = _mtp_backend( _speculative_type = None, + _requested_spec_mode = None, _extra_args = ["--spec-type", "none"], ) assert ( @@ -389,12 +404,17 @@ def test_already_in_target_state_vision_mtp_default_matches(): ) -def test_already_in_target_state_vision_non_mtp_unaffected(): - # Vision non-MTP repo (no -MTP marker) must still mismatch req=None - # against a backend running draft-mtp. +def test_already_in_target_state_vision_off_matches_vision_backend(): + # Vision loads silently drop speculative decoding at the route level + # (_request_matches_loaded_settings overrides req to "off"). At the + # llama_cpp.py level, _already_in_target_state compares canonical + # requested modes; a vision backend recorded with _requested_spec_mode + # = "off" matches a req of "off" or None+vision. backend = _mtp_backend( _model_identifier = "unsloth/Qwen3-VL-4B-Instruct-GGUF", _is_vision = True, + _speculative_type = None, + _requested_spec_mode = "off", ) assert ( backend._already_in_target_state( @@ -403,12 +423,12 @@ def test_already_in_target_state_vision_non_mtp_unaffected(): hf_variant = "Q4_K_M", n_ctx = 8192, cache_type_kv = None, - speculative_type = None, + speculative_type = "off", chat_template_override = None, extra_args = None, is_vision = True, ) - is False + is True ) @@ -482,10 +502,17 @@ def _make_fake_llama_server(path: Path, help_text: str) -> Path: return path +_NEEDS_BASH = pytest.mark.skipif( + sys.platform == "win32", + reason = "fake llama-server is a bash stub; Windows has no direct executor", +) + + def _clear_caps_cache(): LlamaCppBackend._capability_cache.clear() +@_NEEDS_BASH def test_probe_server_capabilities_detects_draft_mtp(tmp_path): # Original naming from llama.cpp #22673. fake = _make_fake_llama_server( @@ -500,6 +527,7 @@ def test_probe_server_capabilities_detects_draft_mtp(tmp_path): assert caps["supports_mtp"] is True +@_NEEDS_BASH def test_probe_server_capabilities_detects_renamed_mtp(tmp_path): # Renamed upstream: draft-mtp -> mtp. fake = _make_fake_llama_server( @@ -513,6 +541,7 @@ def test_probe_server_capabilities_detects_renamed_mtp(tmp_path): assert caps["supports_mtp"] is True +@_NEEDS_BASH def test_probe_server_capabilities_reports_outdated_binary(tmp_path): # Pre-MTP llama.cpp: only ngram variants. fake = _make_fake_llama_server( @@ -533,6 +562,130 @@ def test_probe_server_capabilities_handles_missing_binary(): assert caps["supports_mtp"] is False +# ngram-mod flag flavor detection (new vs legacy llama-server). + +# Help-text fixtures mirror the actual `llama-server --help` block +# layout (flag on its own line; description indented underneath). +_POST_RENAME_HELP = """\ +--spec-draft-n-max N number of tokens to draft for speculative decoding (default: 16) + (env: LLAMA_ARG_SPEC_DRAFT_N_MAX) +--spec-draft-n-min N minimum number of draft tokens to use for speculative decoding (default: 0) + (env: LLAMA_ARG_SPEC_DRAFT_N_MIN) +--spec-draft-p-min, --draft-p-min P minimum speculative decoding probability (greedy) (default: 0.75) + (env: LLAMA_ARG_SPEC_DRAFT_P_MIN) +--spec-ngram-mod-n-min N minimum number of ngram tokens (default: 48) +--spec-ngram-mod-n-max N maximum number of ngram tokens (default: 64) +--spec-ngram-mod-n-match N ngram-mod lookup length (default: 24) +--spec-type none,draft-simple,draft-mtp,ngram-mod comma-separated list of types of speculative decoding to use + (env: LLAMA_ARG_SPEC_TYPE) +--draft, --draft-n, --draft-max N the argument has been removed. use --spec-draft-n-max or --spec-ngram-mod-n-max + (env: LLAMA_ARG_DRAFT_MAX) +--draft-min, --draft-n-min N the argument has been removed. use --spec-draft-n-min or --spec-ngram-mod-n-min + (env: LLAMA_ARG_DRAFT_MIN) +--spec-ngram-size-n N the argument has been removed. use the respective --spec-ngram-*-size-n or --spec-ngram-mod-n-match +""" + +_LEGACY_HELP = """\ +--draft, --draft-n, --draft-max N number of tokens to draft for speculative decoding (default: 8) + (env: LLAMA_ARG_DRAFT_MAX) +--draft-min, --draft-n-min N minimum number of draft tokens to use for speculative decoding (default: 0) + (env: LLAMA_ARG_DRAFT_MIN) +--spec-ngram-size-n N ngram lookup length (default: 24) +--spec-type none,ngram-mod,ngram-simple comma-separated list of types of speculative decoding to use +""" + + +@_NEEDS_BASH +def test_probe_detects_post_rename_ngram_mod_flavor(tmp_path): + fake = _make_fake_llama_server(tmp_path / "llama-server", _POST_RENAME_HELP) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["ngram_mod_flavor"] == "new" + assert caps["supports_ngram_mod"] is True + assert caps["spec_draft_n_max_flag"] == "--spec-draft-n-max" + + +@_NEEDS_BASH +def test_probe_detects_legacy_ngram_mod_flavor(tmp_path): + fake = _make_fake_llama_server(tmp_path / "llama-server", _LEGACY_HELP) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["ngram_mod_flavor"] == "legacy" + assert caps["supports_ngram_mod"] is True + assert caps["spec_draft_n_max_flag"] == "--draft-max" + + +@_NEEDS_BASH +def test_probe_ignores_removal_stub_descriptions(tmp_path): + # Post-rename binary: legacy flags are present but with + # "argument has been removed" descriptions; must not be detected + # as legacy. + fake = _make_fake_llama_server(tmp_path / "llama-server", _POST_RENAME_HELP) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["ngram_mod_flavor"] == "new" + + +@_NEEDS_BASH +def test_probe_no_ngram_mod_on_minimal_binary(tmp_path): + # Pre-anything: neither set present. + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--spec-type none\n--threads N\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["ngram_mod_flavor"] is None + assert caps["supports_ngram_mod"] is False + + +def test_build_ngram_mod_flags_new(): + flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"}) + assert flags == [ + "--spec-ngram-mod-n-match", + "24", + "--spec-ngram-mod-n-min", + "48", + "--spec-ngram-mod-n-max", + "64", + ] + + +def test_build_ngram_mod_flags_legacy(): + flags = _build_ngram_mod_flags({"ngram_mod_flavor": "legacy"}) + assert flags == [ + "--spec-ngram-size-n", + "24", + "--draft-min", + "48", + "--draft-max", + "64", + ] + + +def test_build_ngram_mod_flags_empty_when_unsupported(): + assert _build_ngram_mod_flags({"ngram_mod_flavor": None}) == [] + assert _build_ngram_mod_flags(None) == [] + assert _build_ngram_mod_flags({}) == [] + + +def test_build_ngram_mod_flags_respects_custom_values(): + flags = _build_ngram_mod_flags( + {"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32 + ) + assert flags == [ + "--spec-ngram-mod-n-match", + "16", + "--spec-ngram-mod-n-min", + "24", + "--spec-ngram-mod-n-max", + "32", + ] + + +@_NEEDS_BASH def test_probe_server_capabilities_caches_by_mtime(tmp_path): # Same (path, mtime) -> cache hit. Bumped mtime -> re-probe. fake = _make_fake_llama_server( @@ -555,3 +708,493 @@ def test_probe_server_capabilities_caches_by_mtime(tmp_path): caps2 = LlamaCppBackend.probe_server_capabilities(str(fake)) assert caps2["mtp_token"] == "draft-mtp" assert caps2["supports_mtp"] is True + + +# spec_draft_n_max plumbing (first-class --spec-draft-n-max override). + + +def test_already_in_target_state_matches_when_draft_n_max_unset(): + # None on the request means "platform default"; matches any backend. + backend = _mtp_backend(_spec_draft_n_max = None) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + spec_draft_n_max = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_matches_when_draft_n_max_equals_backend(): + backend = _mtp_backend(_spec_draft_n_max = 4) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + spec_draft_n_max = 4, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_mismatches_when_draft_n_max_differs(): + backend = _mtp_backend(_spec_draft_n_max = 4) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + spec_draft_n_max = 8, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is False + ) + + +def test_already_in_target_state_draft_n_max_ignored_when_not_mtp(): + # ngram-mod backend; spec_draft_n_max is MTP-only and must not force + # a reload against a non-MTP active spec. + backend = _mtp_backend( + _speculative_type = "ngram-mod", + _requested_spec_mode = "ngram", + _spec_draft_n_max = None, + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "ngram-mod", + spec_draft_n_max = 8, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +# Sub-3B MTP gate -- tiny dense models regress with the MTP draft +# head, so load_model falls back to ngram-mod (when the binary supports +# it) instead of draft-mtp. The reload-skip mirror must follow the +# same fallback so a sub-3B reload-with-default does not bounce a +# correctly-configured ngram-mod / off backend. + + +def _patch_probe(monkeypatch, ngram_supported): + """Force probe_server_capabilities to a deterministic result so + tests don't depend on whatever llama-server happens to be on PATH.""" + fake = { + "found": True, + "mtp_token": "draft-mtp", + "supports_mtp": True, + "ngram_mod_flavor": "new" if ngram_supported else None, + "supports_ngram_mod": bool(ngram_supported), + "spec_draft_n_max_flag": "--spec-draft-n-max", + } + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: fake), + ) + monkeypatch.setattr( + LlamaCppBackend, + "_find_llama_server_binary", + classmethod(lambda cls: "/fake/llama-server"), + ) + + +def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported( + monkeypatch, +): + # 0.8B MTP request -- load_model would have promoted to ngram-mod + # (no MTP head); reload check must match a ngram-mod backend. + _patch_probe(monkeypatch, ngram_supported = True) + backend = _mtp_backend( + _model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF", + _speculative_type = "ngram-mod", + _spec_draft_n_max = None, + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_sub_3b_falls_back_to_off_when_no_ngram(monkeypatch): + # 0.8B + binary lacks ngram-mod -> fall back to off. + _patch_probe(monkeypatch, ngram_supported = False) + backend = _mtp_backend( + _model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF", + _speculative_type = None, + _spec_draft_n_max = None, + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_4b_mtp_request_promotes_as_before(monkeypatch): + # 4B is above the 3B threshold -> auto-promote still applies. + _patch_probe(monkeypatch, ngram_supported = True) + backend = _mtp_backend( + _model_identifier = "unsloth/Qwen3.5-4B-MTP-GGUF", + _speculative_type = "draft-mtp", + _spec_draft_n_max = None, + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.5-4B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_2b_falls_back_to_ngram_below_threshold(monkeypatch): + # 2.0B is below the 3B threshold -> ngram-mod fallback, not + # draft-mtp. Clean-bench shows 2B regresses with draft-mtp. + _patch_probe(monkeypatch, ngram_supported = True) + backend = _mtp_backend( + _model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", + _speculative_type = "ngram-mod", + _spec_draft_n_max = None, + ) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +# usage backfill from timings (Studio UI t/s widget fix). + + +def test_backfill_usage_from_timings_fills_when_completion_tokens_zero(): + out = _backfill_usage_from_timings( + {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + {"prompt_n": 42, "predicted_n": 128, "predicted_per_second": 100.0}, + ) + assert out["completion_tokens"] == 128 + assert out["prompt_tokens"] == 42 + assert out["total_tokens"] == 170 + + +def test_backfill_usage_from_timings_fills_when_usage_missing(): + out = _backfill_usage_from_timings( + None, + {"prompt_n": 42, "predicted_n": 128, "predicted_per_second": 100.0}, + ) + assert out["completion_tokens"] == 128 + assert out["prompt_tokens"] == 42 + assert out["total_tokens"] == 170 + + +def test_backfill_usage_from_timings_preserves_real_usage(): + # Non-zero completion_tokens means llama-server reported correctly; + # do not overwrite. + real = {"prompt_tokens": 50, "completion_tokens": 200, "total_tokens": 250} + out = _backfill_usage_from_timings(real, {"predicted_n": 999, "prompt_n": 999}) + assert out is real + assert out["completion_tokens"] == 200 + + +def test_backfill_usage_from_timings_passthrough_when_timings_empty(): + assert _backfill_usage_from_timings(None, None) is None + assert _backfill_usage_from_timings(None, {}) is None + usage = {"completion_tokens": 0} + # No timings.predicted_n -> nothing to fill, return as-is. + assert _backfill_usage_from_timings(usage, {"prompt_ms": 5.0}) is usage + + +# ── _canonicalize_spec_mode (pure) ───────────────────────────────── + + +@pytest.mark.parametrize( + "value, expected", + [ + # New canonical values pass through unchanged. + ("auto", "auto"), + ("mtp", "mtp"), + ("ngram", "ngram"), + ("mtp+ngram", "mtp+ngram"), + ("off", "off"), + ("ngram-simple", "ngram-simple"), + # Legacy wire values map onto the new vocabulary. + ("default", "auto"), + ("draft-mtp", "mtp"), + ("ngram-mod", "ngram"), + # Comma-chained legacy values (e.g. from persisted state) collapse + # to the right canonical mode. + ("ngram-mod,draft-mtp", "mtp+ngram"), + ("draft-mtp,ngram-mod", "mtp+ngram"), + ("draft-mtp,mtp", "mtp"), + ("ngram-mod,ngram", "ngram"), + # Case and whitespace are ignored. + (" AUTO ", "auto"), + ("MTP", "mtp"), + ("MTP+Ngram", "mtp+ngram"), + # None / empty / whitespace pass through as None. + (None, None), + ("", None), + (" ", None), + # Non-string inputs collapse to None. + (42, None), + (True, None), + # Unknown strings fall back to "auto" (safe default). + ("bogus", "auto"), + ], +) +def test_canonicalize_spec_mode(value, expected): + assert _canonicalize_spec_mode(value) == expected + + +# ── _build_speculative_flags resolver matrix ────────────────────── + + +def _resolver_backend(monkeypatch, *, ngram_supported = True, mtp_token = "draft-mtp"): + """Backend with a deterministic probe so the resolver is hermetic.""" + fake = { + "found": True, + "mtp_token": mtp_token, + "supports_mtp": bool(mtp_token), + "ngram_mod_flavor": "new" if ngram_supported else None, + "supports_ngram_mod": bool(ngram_supported), + "spec_draft_n_max_flag": "--spec-draft-n-max", + } + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: fake), + ) + backend = LlamaCppBackend() + backend._nextn_predict_layers = None + return backend + + +def _flags_dict(flags): + """Parse the spec-flag list into a small {flag: value} dict; collapses + repeated flags by keeping the last (only --spec-type can repeat and + never does in our resolver).""" + out = {} + i = 0 + while i < len(flags): + token = flags[i] + if i + 1 < len(flags) and not flags[i + 1].startswith("--"): + out[token] = flags[i + 1] + i += 2 + else: + out[token] = True + i += 1 + return out + + +_MTP_MODEL = "unsloth/Qwen3.6-27B-MTP-GGUF" +_NON_MTP_MODEL = "unsloth/Qwen3-7B-Instruct-GGUF" +_SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF" + + +@pytest.mark.parametrize( + "requested, gpus, model, expect_spec_type, expect_n_max, expect_ngram_knobs", + [ + # ── auto + MTP model + 3B+: GPU = mtp only, CPU = chain ── + ("auto", True, _MTP_MODEL, "draft-mtp", "2", False), + ("auto", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True), + # ── auto + non-MTP: emit --spec-default ── + ("auto", True, _NON_MTP_MODEL, None, None, False), + ("auto", False, _NON_MTP_MODEL, None, None, False), + # ── auto + sub-3B MTP: fallback to ngram-mod ── + ("auto", True, _SUB_3B_MTP_MODEL, "ngram-mod", None, True), + ("auto", False, _SUB_3B_MTP_MODEL, "ngram-mod", None, True), + # ── mtp forced: MTP-only on BOTH platforms ── + ("mtp", True, _MTP_MODEL, "draft-mtp", "2", False), + ("mtp", False, _MTP_MODEL, "draft-mtp", "3", False), + # ── mtp forced on sub-3B: engage anyway ── + ("mtp", True, _SUB_3B_MTP_MODEL, "draft-mtp", "2", False), + # ── mtp forced on non-MTP: engage anyway ── + ("mtp", True, _NON_MTP_MODEL, "draft-mtp", "2", False), + # ── ngram forced: ngram-mod alone on BOTH platforms ── + ("ngram", True, _MTP_MODEL, "ngram-mod", None, True), + ("ngram", False, _MTP_MODEL, "ngram-mod", None, True), + ("ngram", True, _NON_MTP_MODEL, "ngram-mod", None, True), + # ── mtp+ngram forced: chain on BOTH platforms ── + ("mtp+ngram", True, _MTP_MODEL, "ngram-mod,draft-mtp", "2", True), + ("mtp+ngram", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True), + ("mtp+ngram", True, _SUB_3B_MTP_MODEL, "ngram-mod,draft-mtp", "2", True), + # ── off: nothing emitted ── + ("off", True, _MTP_MODEL, None, None, False), + ("off", False, _MTP_MODEL, None, None, False), + # ── legacy values round-trip to the canonical emission ── + ("default", True, _MTP_MODEL, "draft-mtp", "2", False), + ("draft-mtp", True, _MTP_MODEL, "draft-mtp", "2", False), + ("ngram-mod", True, _MTP_MODEL, "ngram-mod", None, True), + ("ngram-mod,draft-mtp", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True), + # ── ngram-simple: pass through ── + ("ngram-simple", True, _MTP_MODEL, "ngram-simple", None, False), + ], +) +def test_build_speculative_flags_matrix( + monkeypatch, + requested, + gpus, + model, + expect_spec_type, + expect_n_max, + expect_ngram_knobs, +): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = requested, + spec_draft_n_max = None, + extra_args = None, + model_identifier = model, + model_path = None, + gpus = gpus, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + if expect_spec_type is None: + assert "--spec-type" not in parsed + else: + assert parsed.get("--spec-type") == expect_spec_type + if expect_n_max is None: + assert "--spec-draft-n-max" not in parsed + else: + assert parsed.get("--spec-draft-n-max") == expect_n_max + if expect_ngram_knobs: + assert "--spec-ngram-mod-n-match" in parsed + assert "--spec-ngram-mod-n-min" in parsed + assert "--spec-ngram-mod-n-max" in parsed + else: + assert "--spec-ngram-mod-n-match" not in parsed + + +def test_build_speculative_flags_user_extra_args_owns_spec_type(monkeypatch): + # User --spec-type in extra_args bypasses the dropdown entirely. + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "mtp", # would normally force MTP + spec_draft_n_max = None, + extra_args = ["--spec-type", "ngram-mod"], + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + # No flags emitted by the resolver -- the user's extra_args carries + # the --spec-type, and the resolver records requested_spec_mode = None. + assert flags == [] + assert backend.requested_spec_mode is None + assert backend.speculative_type is None + + +@pytest.mark.parametrize("mode", ["auto", "mtp", "ngram", "mtp+ngram", "off"]) +def test_build_speculative_flags_round_trips_requested_mode(monkeypatch, mode): + # The status round-trip is the contract that lets the UI dropdown + # restore its picked value after reload / refresh. + backend = _resolver_backend(monkeypatch) + backend._build_speculative_flags( + speculative_type = mode, + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.requested_spec_mode == mode + + +def test_build_speculative_flags_user_draft_n_max_override(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = 5, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-draft-n-max") == "5" + assert backend.spec_draft_n_max == 5 + + +def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch): + # Outdated llama-server with no MTP support: forced MTP must degrade + # to spec-off (warned) rather than emit a bad --spec-type. + backend = _resolver_backend(monkeypatch, mtp_token = None) + flags = backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert "--spec-type" not in flags + # _speculative_type stays None (resolved emission was none), but + # _requested_spec_mode still reflects the user's choice. + assert backend.requested_spec_mode == "mtp" + assert backend.speculative_type is None diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 3703beff0a..1ed0a88b63 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -570,6 +570,11 @@ export function ChatSettingsPanel({ const loadedSpeculativeType = useChatRuntimeStore( (s) => s.loadedSpeculativeType, ); + const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax); + const setSpecDraftNMax = useChatRuntimeStore((s) => s.setSpecDraftNMax); + const loadedSpecDraftNMax = useChatRuntimeStore( + (s) => s.loadedSpecDraftNMax, + ); const modelRequiresTrustRemoteCode = useChatRuntimeStore( (s) => s.modelRequiresTrustRemoteCode, ); @@ -608,7 +613,8 @@ export function ChatSettingsPanel({ const kvDirty = kvCacheDtype !== loadedKvCacheDtype; const ctxDirty = customContextLength !== null; const specDirty = speculativeType !== loadedSpeculativeType; - const modelSettingsDirty = kvDirty || ctxDirty || specDirty; + const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax; + const modelSettingsDirty = kvDirty || ctxDirty || specDirty || specDraftDirty; const chatTemplateOverride = useChatRuntimeStore( (s) => s.chatTemplateOverride, ); @@ -985,19 +991,81 @@ export function ChatSettingsPanel({ Speculative Decoding - Faster generation with 0% accuracy hit. + Faster generation with 0% accuracy hit. Auto picks + MTP / ngram-mod based on the model and platform. + Pick MTP, Ngram, or MTP+Ngram to force a specific + strategy on both GPU and CPU. - { - setSpeculativeType(checked ? "default" : "off"); - }} - /> +
+ +
+ {(speculativeType === "mtp" || + speculativeType === "mtp+ngram") && ( +
+
+ + Draft Tokens + + + Max MTP draft tokens per step + (--spec-draft-n-max). Lower = less wasted + draft decode; higher = bigger speedup when + acceptance stays high. Default: 2 on GPU, + 3 on CPU/Mac. + +
+ { + const raw = e.target.value; + if (raw === "") { + setSpecDraftNMax(null); + return; + } + const parsed = Number.parseInt(raw, 10); + if (Number.isFinite(parsed)) { + const clamped = Math.max(1, Math.min(16, parsed)); + setSpecDraftNMax(clamped); + } + }} + data-test-id="spec-draft-n-max-input" + aria-label="Speculative decoding draft tokens" + className="h-7 w-[72px] rounded-[10px] border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.07] px-2 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0" + /> +
+ )} )} {!isGguf && params.checkpoint && ( @@ -1051,6 +1119,7 @@ export function ChatSettingsPanel({ setCustomContextLength(null); setKvCacheDtype(loadedKvCacheDtype); setSpeculativeType(loadedSpeculativeType); + setSpecDraftNMax(loadedSpecDraftNMax); setChatTemplateOverride(loadedChatTemplateOverride); }} className="h-7 px-3 text-[12px] font-medium tracking-nav text-muted-foreground" 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 03f80c19f2..3f1060edf7 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 @@ -141,10 +141,30 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`; } +// Canonicalises any value the backend reports (or persisted state holds) +// onto the five UI-facing modes the Speculative Decoding dropdown +// understands: "auto" / "mtp" / "ngram" / "mtp+ngram" / "off" / null. +// Mirrors backend _canonicalize_spec_mode so old persisted "default" / +// "draft-mtp" / "ngram-mod" / chain values round-trip cleanly. function normalizeSpeculativeType(v: string | null | undefined): string | null { if (v == null) return null; - if (v === "default" || v === "off") return v; - return "default"; + const s = String(v).trim().toLowerCase(); + if (!s) return null; + if (s === "auto" || s === "default") return "auto"; + if (s === "off") return "off"; + if (s === "ngram-simple") return "ngram-simple"; + if (s === "mtp" || s === "draft-mtp") return "mtp"; + if (s === "ngram" || s === "ngram-mod") return "ngram"; + if (s === "mtp+ngram") return "mtp+ngram"; + // Comma-chained legacy values (e.g. from older persisted state). + const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); + const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod"); + if (hasMtp && hasNgram) return "mtp+ngram"; + if (hasMtp) return "mtp"; + if (hasNgram) return "ngram"; + // Unknown -> safe fallback to Auto so the dropdown stays controlled. + return "auto"; } type LocalReasoningEffort = Extract; @@ -323,6 +343,12 @@ export function useChatModelRuntime() { speculativeType: currentSpecType, loadedSpeculativeType: currentSpecType, }), + ...(statusRes.spec_draft_n_max !== undefined && + prevState.loadedSpecDraftNMax === null && + prevState.specDraftNMax === null && { + specDraftNMax: statusRes.spec_draft_n_max ?? null, + loadedSpecDraftNMax: statusRes.spec_draft_n_max ?? null, + }), ...(statusRes.cache_type_kv !== undefined && prevState.loadedKvCacheDtype === null && { kvCacheDtype: statusRes.cache_type_kv, @@ -528,12 +554,31 @@ export function useChatModelRuntime() { } if (abortCtrl.signal.aborted) throw new Error("Cancelled"); + // Reset Speculative Decoding to Auto whenever the user + // switches to a different model. Spec strategy is a + // per-model decision: a sub-3B non-MTP GGUF that ran with + // "Off" should not carry that choice into a 27B MTP GGUF + // where Auto would auto-promote to draft-mtp. The user can + // still pick a forced mode on the new model; this just + // clears the stale prior-model choice so the backend's + // platform-aware path runs by default. Same applies to + // spec_draft_n_max which is MTP-only. + if (currentCheckpoint && currentCheckpoint !== modelId) { + useChatRuntimeStore.setState({ + speculativeType: null, + loadedSpeculativeType: null, + specDraftNMax: null, + loadedSpecDraftNMax: null, + }); + } + const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength, speculativeType, + specDraftNMax, activePresetSource, activeGgufVariant, } = useChatRuntimeStore.getState(); @@ -561,6 +606,7 @@ export function useChatModelRuntime() { chat_template_override: effectiveChatTemplateOverride, cache_type_kv: kvCacheDtype, speculative_type: speculativeType, + spec_draft_n_max: specDraftNMax, }); // If cancelled while loading, don't update UI to show @@ -635,6 +681,8 @@ export function useChatModelRuntime() { loadedKvCacheDtype: loadedKv, speculativeType: loadedSpec, loadedSpeculativeType: loadedSpec, + specDraftNMax: loadResponse.spec_draft_n_max ?? null, + loadedSpecDraftNMax: loadResponse.spec_draft_n_max ?? null, customContextLength: keepCustomCtx, defaultChatTemplate: loadResponse.chat_template ?? null, chatTemplateOverride: effectiveChatTemplateOverride, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 13d3cb533f..e55a25d08d 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -258,6 +258,9 @@ type ChatRuntimeStore = { loadedKvCacheDtype: string | null; speculativeType: string | null; loadedSpeculativeType: string | null; + /** User --spec-draft-n-max override (null = platform default). */ + specDraftNMax: number | null; + loadedSpecDraftNMax: number | null; loadedIsMultimodal: boolean; customContextLength: number | null; defaultChatTemplate: string | null; @@ -305,6 +308,7 @@ type ChatRuntimeStore = { setToolCallTimeout: (value: number) => void; setKvCacheDtype: (dtype: string | null) => void; setSpeculativeType: (type: string | null) => void; + setSpecDraftNMax: (value: number | null) => void; setCustomContextLength: (v: number | null) => void; setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; @@ -349,8 +353,10 @@ export const useChatRuntimeStore = create((set) => ({ toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5), kvCacheDtype: null, loadedKvCacheDtype: null, - speculativeType: "default", + speculativeType: "auto", loadedSpeculativeType: null, + specDraftNMax: null, + loadedSpecDraftNMax: null, loadedIsMultimodal: false, customContextLength: null, defaultChatTemplate: null, @@ -457,8 +463,10 @@ export const useChatRuntimeStore = create((set) => ({ toolStatus: null, kvCacheDtype: null, loadedKvCacheDtype: null, - speculativeType: "default", + speculativeType: "auto", loadedSpeculativeType: null, + specDraftNMax: null, + loadedSpecDraftNMax: null, loadedIsMultimodal: false, customContextLength: null, defaultChatTemplate: null, @@ -506,6 +514,7 @@ export const useChatRuntimeStore = create((set) => ({ }), setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }), setSpeculativeType: (speculativeType) => set({ speculativeType }), + setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }), setCustomContextLength: (customContextLength) => set({ customContextLength }), setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }), setPendingAudio: (base64, name) => diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 3fc5d320df..1e6bcf8b87 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -42,7 +42,21 @@ export interface LoadModelRequest { trust_remote_code?: boolean; chat_template_override?: string | null; cache_type_kv?: string | null; + /** + * Speculative decoding mode for GGUF models. Canonical values: + * "auto" (platform-aware: MTP on MTP GGUFs, ngram-mod fallback for + * sub-3B), "mtp" (force draft-mtp only on both GPU and CPU), "ngram" + * (force ngram-mod only), "mtp+ngram" (force ngram-mod + draft-mtp + * chain on both platforms), or "off". Legacy values "default" / + * "draft-mtp" / "ngram-mod" / "ngram-simple" are still accepted by + * the backend. + */ speculative_type?: string | null; + /** + * Override --spec-draft-n-max for MTP speculative decoding. Only + * applied when speculative_type resolves to "mtp" or "mtp+ngram". + */ + spec_draft_n_max?: number | null; } export interface ValidateModelResponse { @@ -118,7 +132,9 @@ export interface LoadModelResponse { supports_tools?: boolean; cache_type_kv?: string | null; chat_template?: string | null; + /** Canonical UI-facing mode the load request resolved to. See LoadModelRequest. */ speculative_type?: string | null; + spec_draft_n_max?: number | null; } export interface UnloadModelRequest { @@ -155,7 +171,9 @@ export interface InferenceStatusResponse { native_context_length?: number | null; cache_type_kv?: string | null; chat_template_override?: string | null; + /** Canonical UI-facing mode currently active. See LoadModelRequest. */ speculative_type?: string | null; + spec_draft_n_max?: number | null; } export interface AudioGenerationResponse { From bef6da59aac7e14a92e46bb6537fdced9072a7e6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:19:02 -0700 Subject: [PATCH 18/32] studio: reserve VRAM headroom for the MTP draft cache in auto-fit (#5585) * studio: reserve VRAM headroom for the MTP draft cache in auto-fit When MTP is going to engage on this load, _fit_context_to_vram now budgets 0.85 of available VRAM instead of 0.90, leaving room for llama.cpp's secondary MTP draft KV cache + compute graph buffers. Motivation: a user report on RTX 5090 (32 GB) showed Qwen3.6-27B-MTP-GGUF UD-Q4_K_XL at native auto-context running roughly half the speed of the same model with a slightly smaller context. The most parsimonious explanation is a VRAM cliff: at native context the target's KV already eats the 90% budget, then llama-server allocates the draft cache + draft graph on top and spills into a slower partial-offload path. Reducing the budget by 5% on MTP loads avoids the spill without penalising non-MTP loads. On hardware with abundant VRAM (B200, etc.) the fit is unchanged because the requested context already fits in the tighter budget too. MTP detection mirrors the auto-promotion logic in load_model: the GGUF advertises nextn_predict_layers, or the model identifier / local path matches the -MTP marker, and the user has not explicitly opted out via speculative_type="off" or --spec-type extra args. Tests: two new cases in test_kv_cache_estimation.py verify that mtp_engaged=True yields a context less-than-or-equal-to the non-MTP path on a tight budget, and that kv_on_gpu=False still short-circuits regardless of mtp_engaged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: gate _mtp_will_engage on canonical-mode resolver After PR #5582 introduced the 5-mode Speculative Decoding dropdown plus _canonicalize_spec_mode, the auto-fit MTP-engaged predicate becomes: * forced mtp / mtp+ngram -> always engage MTP (extra VRAM needed) * auto + MTP GGUF (>= 3B) -> engages MTP via auto-promotion * auto + MTP GGUF (sub-3B) -> falls back to ngram-mod (no extra VRAM) * ngram / ngram-simple / off -> never engage MTP * user --spec-type in extra_args -> resolver suppressed; no headroom The old gate triggered on "anything but off", so it over-reserved the 0.85 budget when the user explicitly picked Ngram (no MTP) or when Auto fell back to ngram-mod on a sub-3B MTP model. The 5% headroom cost was minor but unnecessary. Mirrors the same logic already encoded in _build_speculative_flags so the auto-fit budget and the actual emission agree on whether MTP is running. All 361 backend tests pass. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 42 ++++++++++++++++++- .../backend/tests/test_kv_cache_estimation.py | 27 ++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 134b6574b0..91ee1e4203 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1695,6 +1695,7 @@ class LlamaCppBackend: kv_unified: bool = True, ctx_checkpoints: int = 0, kv_on_gpu: bool = True, + mtp_engaged: bool = False, ) -> int: """Return the largest context length that fits in GPU VRAM. @@ -1708,6 +1709,12 @@ class LlamaCppBackend: the KV cache lives in CPU RAM and doesn't compete with weights for VRAM; the requested context is honored verbatim. The other keyword args mirror ``_estimate_kv_cache_bytes``. + + ``mtp_engaged`` reserves extra VRAM for the MTP draft model's + KV cache + compute graph buffers. llama.cpp's MTP path keeps a + secondary cache sized off the target's KV; on tight VRAM tiers + (e.g. 32 GB) auto-fit at native context would otherwise spill + and force llama-server into a slower partial-offload path. """ if not self._can_estimate_kv(): logger.debug( @@ -1728,7 +1735,9 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - budget_bytes = available_mib * 1024 * 1024 * 0.90 + # MTP needs a tighter budget; drop from 0.90 to 0.85. + budget_frac = 0.85 if mtp_engaged else 0.90 + budget_bytes = available_mib * 1024 * 1024 * budget_frac model_footprint = model_size_bytes # Check if requested context already fits @@ -2614,6 +2623,35 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink this if hardware is limited. max_available_ctx = self._context_length or effective_ctx + # Will MTP engage on this load? If so, the auto-fit + # budget needs to reserve extra VRAM for the draft + # model's KV cache + compute graph. Mirrors the + # canonical-mode resolver in _build_speculative_flags: + # forced mtp / mtp+ngram always engage; auto only + # engages on an MTP GGUF >= 3B (sub-3B auto falls + # back to ngram-mod which doesn't need headroom); + # ngram / ngram-simple / off never engage MTP. + _mtp_canonical = _canonicalize_spec_mode(speculative_type) + _mtp_effective = _mtp_canonical or "auto" + _mtp_size_for_fit = _extract_model_size_b(model_identifier) + _mtp_sub_3b_for_fit = ( + _mtp_size_for_fit is not None and _mtp_size_for_fit < 3.0 + ) + _mtp_will_engage = bool( + not _extra_args_set_spec_type(extra_args) + and ( + _mtp_effective in ("mtp", "mtp+ngram") + or ( + _mtp_effective == "auto" + and ( + bool(self._nextn_predict_layers) + or _is_mtp_model_name(model_identifier, model_path) + ) + and not _mtp_sub_3b_for_fit + ) + ) + ) + # Auto-cap context to fit in GPU VRAM and select GPUs. # # Two policies depending on whether the user set n_ctx: @@ -2649,6 +2687,7 @@ class LlamaCppBackend: model_size, cache_type_kv, n_parallel = n_parallel, + mtp_engaged = _mtp_will_engage, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel @@ -2700,6 +2739,7 @@ class LlamaCppBackend: model_size, cache_type_kv, n_parallel = n_parallel, + mtp_engaged = _mtp_will_engage, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 29d87804ff..d52a58a25c 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -1558,6 +1558,33 @@ class TestServerFlags: ) assert fitted < 32_768 + def test_fit_mtp_engaged_returns_smaller_or_equal_context(self): + # MTP-engaged budget is 0.85 of available; non-MTP is 0.90. + # On a tight budget the MTP path must yield <= the non-MTP path. + b = self._gqa_backend() + common = dict( + requested_ctx = 32_768, + available_mib = 128, + model_size_bytes = 8 * 1024 * 1024, + cache_type_kv = "f16", + ) + baseline = b._fit_context_to_vram(**common) + mtp = b._fit_context_to_vram(**common, mtp_engaged = True) + assert mtp <= baseline + + def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self): + # kv_on_gpu=False short-circuits the fit; mtp_engaged is irrelevant. + b = self._gqa_backend() + fitted = b._fit_context_to_vram( + requested_ctx = 32_768, + available_mib = 1, + model_size_bytes = 100, + cache_type_kv = "f16", + kv_on_gpu = False, + mtp_engaged = True, + ) + assert fitted == 32_768 + def test_fit_threads_swa_full_through_estimator(self): # SWA model, generous budget; both should fit but cache size differs. b = self._swa_backend() From bb4eb88fdcc7e156b3122be9c4a902e7c2488594 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:30:17 -0700 Subject: [PATCH 19/32] Studio: tools, thinking blocks, code execution and web search for safetensors (#5520) Adds tools, thinking blocks, code execution, and web search support to the safetensors / transformers and MLX inference backends in Studio, bringing them to parity with the GGUF path. What ships - safetensors / transformers agentic tool loop with cumulative-text state machine, tool-call XML parser, and template kwarg forwarding (tools / enable_thinking / reasoning_effort / preserve_thinking). - MLX backend: same kwargs accepted on Apple Silicon; chat_template_info shipped through worker IPC; pills enable for Qwen / Qwen3 / Qwen3.5 / Gemma reasoning. - Capability classifier (_detect_safetensors_features) gates supports_tools on actual parser-compatible emission markers ( / XML). - CWE-209 hygiene: safetensors SSE error path emits a constant message and logs the trace server-side. Validation - 256 unit tests green (43 tool-loop, 11 capability advertise, 7 MLX backend, 5 main-added, 190 adjacent inference / anthropic / openai regression). - Cross-OS staging CI green on ubuntu-latest / macos-14 / windows-latest plus a dedicated MLX cartesian probe against real unsloth/Qwen3.5-0.8B on macos-14 (CI 26098107440). - Capability parity verified across Qwen3 / Qwen3.5 / Llama-3 / Mistral / Gemma / DeepSeek-R1 / gpt-oss (incl. BF16). - Manual confirmation from Imagineer99 on Qwen3.5-2B: think + search + code exec working. Closes the safetensors / MLX gap with the GGUF backend. --- .../core/inference/chat_template_helpers.py | 60 ++ studio/backend/core/inference/inference.py | 138 ++- studio/backend/core/inference/llama_cpp.py | 16 +- .../backend/core/inference/mlx_inference.py | 113 ++- studio/backend/core/inference/orchestrator.py | 142 +++- .../core/inference/safetensors_agentic.py | 392 +++++++++ .../core/inference/tool_call_parser.py | 204 +++++ studio/backend/core/inference/worker.py | 62 +- studio/backend/routes/inference.py | 434 ++++++++-- .../tests/test_mlx_inference_backend.py | 94 +++ .../test_safetensors_capability_advertise.py | 451 ++++++++++ .../tests/test_safetensors_tool_loop.py | 788 ++++++++++++++++++ studio/backend/utils/datasets/__init__.py | 2 + .../backend/utils/datasets/model_mappings.py | 20 + .../src/features/chat/api/chat-adapter.ts | 6 + 15 files changed, 2798 insertions(+), 124 deletions(-) create mode 100644 studio/backend/core/inference/chat_template_helpers.py create mode 100644 studio/backend/core/inference/safetensors_agentic.py create mode 100644 studio/backend/core/inference/tool_call_parser.py create mode 100644 studio/backend/tests/test_safetensors_capability_advertise.py create mode 100644 studio/backend/tests/test_safetensors_tool_loop.py diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py new file mode 100644 index 0000000000..833a714ee4 --- /dev/null +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Dependency-light wrapper around tokenizer.apply_chat_template with a +kwarg fallback for templates that reject reasoning/tools args. +""" + +from typing import Optional + + +def apply_chat_template_for_generation( + tokenizer, + messages: list, + *, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, +) -> str: + """Render the chat prompt. Try richest kwargs first; drop one + group at a time on TypeError. Jinja / missing-variable errors + propagate.""" + reasoning_kwargs: dict = {} + if enable_thinking is not None: + reasoning_kwargs["enable_thinking"] = enable_thinking + if reasoning_effort is not None: + reasoning_kwargs["reasoning_effort"] = reasoning_effort + if preserve_thinking is not None: + reasoning_kwargs["preserve_thinking"] = preserve_thinking + + attempts: list[dict] = [] + if tools and reasoning_kwargs: + attempts.append({"tools": tools, **reasoning_kwargs}) + if tools: + attempts.append({"tools": tools}) + if reasoning_kwargs: + attempts.append(dict(reasoning_kwargs)) + attempts.append({}) + + last_exc: Optional[Exception] = None + for kwargs in attempts: + try: + return tokenizer.apply_chat_template( + messages, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ) + except TypeError as e: + last_exc = e + continue + except Exception as e: + last_exc = e + break + if last_exc is not None: + raise last_exc + raise RuntimeError( + "apply_chat_template_for_generation: no attempt produced a result" + ) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4c140013a0..e1620f5ca3 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -839,6 +839,74 @@ class InferenceBackend: cancel_event = cancel_event, _adapter_state = use_adapter, **gen_kwargs ) + def generate_chat_completion_with_tools( + self, + messages: list, + tools: list, + system_prompt: str = "", + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 2048, + repetition_penalty: float = 1.0, + cancel_event = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + max_tool_iterations: int = 25, + auto_heal_tool_calls: bool = True, + tool_call_timeout: int = 300, + session_id: Optional[str] = None, + ): + """Run an agentic tool loop on top of ``generate_chat_response``. + + Yields the same event-dict protocol used by the GGUF path so + the route layer can stream both backends through one helper. + Each event is one of: + + * ``{"type": "status", "text": ...}`` + * ``{"type": "content", "text": cumulative_text}`` + * ``{"type": "tool_start", "tool_name", "tool_call_id", "arguments"}`` + * ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}`` + """ + from core.inference.safetensors_agentic import run_safetensors_tool_loop + from core.inference.tools import execute_tool + + def _single_turn(conv: list): + # conv already has the system message -- avoid double-prepend. + yield from self._generate_chat_response_inner( + messages = conv, + system_prompt = "", + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + cancel_event = cancel_event, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + + initial = list(messages) + if system_prompt: + initial = [{"role": "system", "content": system_prompt}] + initial + + yield from run_safetensors_tool_loop( + single_turn = _single_turn, + messages = initial, + tools = tools, + execute_tool = execute_tool, + cancel_event = cancel_event, + auto_heal_tool_calls = auto_heal_tool_calls, + max_tool_iterations = max_tool_iterations, + tool_call_timeout = tool_call_timeout, + session_id = session_id, + ) + def generate_chat_response( self, messages: list, @@ -851,10 +919,20 @@ class InferenceBackend: max_new_tokens: int = 256, repetition_penalty: float = 1.0, cancel_event = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, ) -> Generator[str, None, None]: """ Generate response for text or vision models. The generation lock is acquired by the background generation thread. + + ``tools`` / ``enable_thinking`` / ``reasoning_effort`` / + ``preserve_thinking`` are forwarded into + ``tokenizer.apply_chat_template`` so templates that understand + these kwargs (Qwen3, Llama 3.1+, gpt-oss harmony, ...) advertise + the tool schemas and reasoning controls to the model. """ yield from self._generate_chat_response_inner( messages = messages, @@ -867,6 +945,10 @@ class InferenceBackend: max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, cancel_event = cancel_event, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, ) def _generate_chat_response_inner( @@ -882,6 +964,10 @@ class InferenceBackend: repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, ) -> Generator[str, None, None]: """ Inner generation logic. Called by both generate_chat_response @@ -981,8 +1067,13 @@ class InferenceBackend: f"Please use a model that includes a chat template, or manually set " f"one via tokenizer.chat_template before inference." ) - formatted_prompt = tokenizer.apply_chat_template( - template_messages, tokenize = False, add_generation_prompt = True + formatted_prompt = self._apply_chat_template_for_generation( + tokenizer, + template_messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, ) logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: @@ -1319,20 +1410,9 @@ class InferenceBackend: def _is_gpt_oss_model(self, model_name: str = None) -> bool: """Check if the given (or active) model uses the gpt-oss harmony protocol.""" - name = (model_name or self.active_model_name or "").lower() - try: - from utils.datasets import MODEL_TO_TEMPLATE_MAPPER + from utils.datasets import is_gpt_oss_model_name - # Exact match - if MODEL_TO_TEMPLATE_MAPPER.get(name) == "gpt-oss": - return True - # Partial match (e.g. name-bnb-4bit variants) - for key, tmpl in MODEL_TO_TEMPLATE_MAPPER.items(): - if tmpl == "gpt-oss" and (key in name or name in key): - return True - except Exception: - pass - return "gpt-oss" in name + return is_gpt_oss_model_name(model_name or self.active_model_name or "") def generate_stream( self, @@ -1715,6 +1795,34 @@ class InferenceBackend: "Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS" ) + def _apply_chat_template_for_generation( + self, + tokenizer, + messages: list, + *, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + ) -> str: + """Render the chat prompt, peeling kwargs the template does not + understand. Delegates to the dependency-light helper module so + the fallback chain can be unit-tested without pulling unsloth / + torch into the test sandbox. + """ + from core.inference.chat_template_helpers import ( + apply_chat_template_for_generation, + ) + + return apply_chat_template_for_generation( + tokenizer, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str: if not self.active_model_name or self.active_model_name not in self.models: logger.error("No active model available") diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 91ee1e4203..260e675a73 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -44,6 +44,9 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from core.inference.tool_call_parser import ( + parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, +) logger = get_logger(__name__) @@ -3904,16 +3907,9 @@ class LlamaCppBackend: @staticmethod def _parse_tool_calls_from_text(content: str) -> list[dict]: - """ - Parse tool calls from XML markup in content text. - - Handles formats like: - {"name":"web_search","arguments":{"query":"..."}} - ... - Closing tags (, , ) are all optional - since models frequently omit them. - """ - return parse_tool_calls_from_text(content) + """Thin wrapper around the shared parser in tool_call_parser + so safetensors and llama_cpp pick up the same fixes.""" + return _shared_parse_tool_calls_from_text(content) @staticmethod def _build_openai_messages( diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index e7bce2d33e..716e4c27a2 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -157,10 +157,59 @@ class MLXInferenceBackend: "audio_type": None, "has_audio_input": False, } + # Capture chat_template_info so the worker IPC reply can ship + # it back to the parent and the route layer classifies + # capabilities the same way as the transformers / GGUF paths. + self._populate_chat_template_info(model_name) logger.info("Model %s loaded successfully", model_name) return True + def _populate_chat_template_info(self, model_name: str) -> None: + """Mirror InferenceBackend._load_chat_template_info for MLX. + + Stores ``chat_template_info`` on ``self.models[model_name]`` + with the resolved ``tokenizer.chat_template`` so + ``_detect_safetensors_features`` (route layer) sees the same + template the model actually uses.""" + entry = self.models.get(model_name) + if not entry: + return + tok = entry.get("tokenizer") + if tok is None: + proc = entry.get("processor") + tok = getattr(proc, "tokenizer", None) if proc else None + info = { + "has_template": False, + "template": None, + "format_type": "generic", + "special_tokens": {}, + "template_name": None, + } + try: + tpl = getattr(tok, "chat_template", None) + if tpl: + info["has_template"] = True + info["template"] = tpl + lower = tpl.lower() + if "start_header_id" in lower and "end_header_id" in lower: + info["format_type"] = "llama3" + elif "[inst]" in lower and "[/inst]" in lower: + info["format_type"] = "mistral" + elif "<|im_start|>" in lower and "<|im_end|>" in lower: + info["format_type"] = "chatml" + else: + info["format_type"] = "custom" + special = {} + for attr in ("bos_token", "eos_token", "pad_token"): + val = getattr(tok, attr, None) + if val: + special[attr] = val + info["special_tokens"] = special + except Exception as exc: + logger.warning("MLX chat_template_info capture failed: %s", exc) + entry["chat_template_info"] = info + def unload_model(self, model_name: str) -> bool: import mlx.core as mx import gc @@ -197,6 +246,14 @@ class MLXInferenceBackend: max_new_tokens = 256, repetition_penalty = 1.0, cancel_event = None, + # Reasoning / tool kwargs forwarded by the route + worker -- the + # MLX path renders the template via apply_chat_template_for_ + # generation so these are honoured the same way as the + # transformers path. + tools = None, + enable_thinking = None, + reasoning_effort = None, + preserve_thinking = None, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -239,6 +296,10 @@ class MLXInferenceBackend: max_new_tokens, repetition_penalty, cancel_event, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, ) else: yield from self._generate_text( @@ -250,6 +311,10 @@ class MLXInferenceBackend: max_new_tokens, repetition_penalty, cancel_event, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, ) def _generate_text( @@ -262,14 +327,26 @@ class MLXInferenceBackend: max_new_tokens, repetition_penalty, cancel_event, + *, + tools = None, + enable_thinking = None, + reasoning_effort = None, + preserve_thinking = None, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors - prompt = self._tokenizer.apply_chat_template( + from core.inference.chat_template_helpers import ( + apply_chat_template_for_generation, + ) + + prompt = apply_chat_template_for_generation( + self._tokenizer, messages, - tokenize = False, - add_generation_prompt = True, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, ) if prompt is None: raise RuntimeError( @@ -343,20 +420,38 @@ class MLXInferenceBackend: max_new_tokens, repetition_penalty, cancel_event, + *, + tools = None, + enable_thinking = None, + reasoning_effort = None, + preserve_thinking = None, ): from mlx_vlm import stream_generate as vlm_stream - # Apply chat template - chat_fn = getattr(self._processor, "apply_chat_template", None) + from core.inference.chat_template_helpers import ( + apply_chat_template_for_generation, + ) + + # Pick the chat-template-aware caller: processors that expose + # their own apply_chat_template + chat_template attr (e.g. + # Qwen2.5-VL) use it directly; otherwise fall back to the + # nested tokenizer. + chat_target = self._processor if ( - chat_fn is None + getattr(self._processor, "apply_chat_template", None) is None or not hasattr(self._processor, "chat_template") or self._processor.chat_template is None ): - tok = getattr(self._processor, "tokenizer", self._processor) - chat_fn = tok.apply_chat_template + chat_target = getattr(self._processor, "tokenizer", self._processor) - prompt = chat_fn(messages, tokenize = False, add_generation_prompt = True) + prompt = apply_chat_template_for_generation( + chat_target, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) # For VLM: always use mlx_vlm's stream_generate which handles # pixel_values properly (passes None for text-only, image for VLM) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 5562820f49..7e7d7026f6 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -449,6 +449,10 @@ class InferenceOrchestrator: repetition_penalty: float = 1.0, cancel_event = None, use_adapter = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, ) -> Generator[str, None, None]: """Dispatched generation — sends command without holding _gen_lock. @@ -494,6 +498,14 @@ class InferenceOrchestrator: if use_adapter is not None: cmd["use_adapter"] = use_adapter + if tools is not None: + cmd["tools"] = tools + if enable_thinking is not None: + cmd["enable_thinking"] = enable_thinking + if reasoning_effort is not None: + cmd["reasoning_effort"] = reasoning_effort + if preserve_thinking is not None: + cmd["preserve_thinking"] = preserve_thinking # Create mailbox BEFORE sending command mailbox: queue.Queue = queue.Queue() @@ -695,6 +707,13 @@ class InferenceOrchestrator: "audio_type": model_info.get("audio_type"), "has_audio_input": model_info.get("has_audio_input", False), } + # Mirror chat_template_info so routes can classify + # capabilities without re-entering the subprocess. + _tpl_info = model_info.get("chat_template_info") + if isinstance(_tpl_info, dict): + self.models[self.active_model_name]["chat_template_info"] = ( + _tpl_info + ) self.loading_models.discard(model_name) logger.info( "Model '%s' loaded successfully in subprocess", model_name @@ -770,8 +789,18 @@ class InferenceOrchestrator: max_new_tokens: int = 256, repetition_penalty: float = 1.0, cancel_event = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, ) -> Generator[str, None, None]: - """Generate response, streaming tokens from subprocess.""" + """Generate response, streaming tokens from subprocess. + + Optional ``tools`` / ``enable_thinking`` / ``reasoning_effort`` / + ``preserve_thinking`` kwargs are forwarded into the worker so + ``tokenizer.apply_chat_template`` can render tool schemas and + reasoning controls when the template understands them. + """ yield from self._generate_inner( messages = messages, system_prompt = system_prompt, @@ -784,6 +813,88 @@ class InferenceOrchestrator: repetition_penalty = repetition_penalty, cancel_event = cancel_event, use_adapter = None, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + + def generate_chat_completion_with_tools( + self, + messages: list, + tools: list, + system_prompt: str = "", + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_tokens: Optional[int] = None, + repetition_penalty: float = 1.0, + cancel_event = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + max_tool_iterations: int = 25, + auto_heal_tool_calls: bool = True, + tool_call_timeout: int = 300, + session_id: Optional[str] = None, + use_adapter: Optional[Union[bool, str]] = None, + **_unused, + ): + """Run the safetensors agentic tool loop in this (parent) + process, calling the worker for each generation turn. + + Yields the same event dicts as the GGUF tool loop so the route + layer can stream both backends through one helper. See + ``safetensors_agentic.run_safetensors_tool_loop`` for the + event protocol. + """ + from core.inference.safetensors_agentic import run_safetensors_tool_loop + from core.inference.tools import execute_tool + + max_new_tokens = max_tokens if max_tokens and max_tokens > 0 else 2048 + + def _single_turn(conv: list): + # ``conv`` already carries any system message because the + # loop appends to a list seeded with system+user above. + common_kwargs = dict( + messages = conv, + system_prompt = "", + image = None, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + cancel_event = cancel_event, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + if use_adapter is not None: + yield from self.generate_with_adapter_control( + use_adapter = use_adapter, + **common_kwargs, + ) + else: + yield from self.generate_chat_response(**common_kwargs) + + initial = list(messages) + if system_prompt: + initial = [{"role": "system", "content": system_prompt}] + initial + + yield from run_safetensors_tool_loop( + single_turn = _single_turn, + messages = initial, + tools = tools, + execute_tool = execute_tool, + cancel_event = cancel_event, + auto_heal_tool_calls = auto_heal_tool_calls, + max_tool_iterations = max_tool_iterations, + tool_call_timeout = tool_call_timeout, + session_id = session_id, ) def generate_with_adapter_control( @@ -817,6 +928,10 @@ class InferenceOrchestrator: repetition_penalty: float = 1.0, cancel_event = None, use_adapter = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, ) -> Generator[str, None, None]: """Inner generation logic — sends command to subprocess, yields tokens. @@ -853,6 +968,10 @@ class InferenceOrchestrator: repetition_penalty = repetition_penalty, cancel_event = cancel_event, use_adapter = use_adapter, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, ) def _generate_locked( @@ -868,6 +987,10 @@ class InferenceOrchestrator: repetition_penalty: float = 1.0, cancel_event = None, use_adapter = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, ) -> Generator[str, None, None]: """Actual generation logic — must be called under _gen_lock.""" request_id = str(uuid.uuid4()) @@ -893,6 +1016,16 @@ class InferenceOrchestrator: if use_adapter is not None: cmd["use_adapter"] = use_adapter + # Only forward template kwargs the caller actually set so older + # workers that ignore unknown keys still work. + if tools is not None: + cmd["tools"] = tools + if enable_thinking is not None: + cmd["enable_thinking"] = enable_thinking + if reasoning_effort is not None: + cmd["reasoning_effort"] = reasoning_effort + if preserve_thinking is not None: + cmd["preserve_thinking"] = preserve_thinking try: self._send_cmd(cmd) @@ -1200,6 +1333,13 @@ class InferenceOrchestrator: return self.models[self.active_model_name].get("is_vision", False) return False + def _is_gpt_oss_model(self, model_name: str = None) -> bool: + """Parent-side gpt-oss detection so the safetensors route can run + the same guard without an IPC round-trip to the subprocess.""" + from utils.datasets import is_gpt_oss_model_name + + return is_gpt_oss_model_name(model_name or self.active_model_name or "") + # ========== GLOBAL INSTANCE ========== _inference_backend = None diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py new file mode 100644 index 0000000000..73bb3d090a --- /dev/null +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Safetensors/transformers agentic tool loop. + +Wraps a single-turn cumulative-text generator (the existing +``InferenceOrchestrator.generate_chat_response`` pipeline that streams +from a worker subprocess) with the tool-calling, thinking-block, +status, and metadata event protocol used by the GGUF path. Keeps the +front-end SSE shape identical across backends so the chat UI does not +care which engine actually ran the model. + +The GGUF path lives in ``llama_cpp.py`` and talks to llama-server's +structured ``delta.tool_calls`` directly. Native transformers has no +such structured channel, so this loop parses tool calls from the +cumulative text and dispatches them via ``core.inference.tools``. +""" + +import json +import threading +from typing import Callable, Generator, Optional +from urllib.parse import urlparse + +from loggers import get_logger + +from core.inference.tool_call_parser import ( + BUDGET_EXHAUSTED_NUDGE, + DUPLICATE_CALL_NUDGE, + TOOL_ERROR_NUDGE, + TOOL_ERROR_PREFIXES, + TOOL_XML_SIGNALS, + has_tool_signal, + parse_tool_calls_from_text, + strip_tool_markup, +) + + +logger = get_logger(__name__) + + +# Buffer cap while waiting to disambiguate a possible tool-call prefix. +_MAX_BUFFER_CHARS = 32 + + +def _status_for_tool(tool_name: str, arguments: dict) -> str: + """Return a human-readable status line matching the GGUF path.""" + if tool_name == "web_search": + url = (arguments.get("url") or "").strip() + if url: + parsed = urlparse(url) + if parsed.scheme in ("http", "https") and parsed.hostname: + host = parsed.hostname + if host.startswith("www."): + host = host[4:] + return f"Reading: {host}" + return "Reading page..." + query = arguments.get("query", "") + return f"Searching: {query}" + if tool_name == "python": + preview = (arguments.get("code") or "").strip().split("\n")[0][:60] + return f"Running Python: {preview}" if preview else "Running Python..." + if tool_name == "terminal": + preview = (arguments.get("command") or "")[:60] + return f"Running: {preview}" if preview else "Running command..." + return f"Calling: {tool_name}" + + +_CANONICAL_HEAL_ARG = {"python": "code", "terminal": "command"} + + +def _coerce_arguments(raw_args, *, heal: bool, tool_name: str = "") -> dict: + """Normalise tool ``arguments`` to a dict. + + Some templates emit a JSON string, others a bare query string. With + ``heal=True`` we accept a bare string as ``{: ...}`` + so a Hermes-style call without proper JSON still runs the tool. The + canonical key is picked per tool: ``code`` for python, ``command`` + for terminal, ``query`` for everything else (e.g. web_search). + """ + if isinstance(raw_args, dict): + return raw_args + if isinstance(raw_args, str): + try: + parsed = json.loads(raw_args) + if isinstance(parsed, dict): + return parsed + except (json.JSONDecodeError, ValueError): + pass + if heal: + key = _CANONICAL_HEAL_ARG.get(tool_name, "query") + return {key: raw_args} + return {"raw": raw_args} + return {} + + +def run_safetensors_tool_loop( + *, + single_turn: Callable[[list], Generator[str, None, None]], + messages: list[dict], + tools: list[dict], + execute_tool: Callable[..., str], + cancel_event: Optional[threading.Event] = None, + auto_heal_tool_calls: bool = True, + max_tool_iterations: int = 25, + tool_call_timeout: int = 300, + session_id: Optional[str] = None, +) -> Generator[dict, None, None]: + """Drive an agentic tool loop on top of a cumulative-text generator. + + ``single_turn(messages)`` must yield cumulative assistant text + (each yield is a snapshot including all previously emitted tokens). + The loop: + + * Buffers the leading characters of every turn so it can decide + whether the model is about to emit a tool call. Plain content + starts streaming as soon as the buffer rules it out. + * On detecting ```` or ``= 0 and (signal_pos < 0 or p < signal_pos): + signal_pos = p + if signal_pos >= 0: + before_tool = candidate[:signal_pos] + cleaned_before = strip_tool_markup(before_tool) + if len(cleaned_before) > len(last_emitted): + last_emitted = cleaned_before + yield {"type": "content", "text": cleaned_before} + cumulative_display = candidate + detect_state = _state_draining + continue + cumulative_display = candidate + cleaned = strip_tool_markup(cumulative_display) + if len(cleaned) > len(last_emitted): + last_emitted = cleaned + yield {"type": "content", "text": cleaned} + continue + + # BUFFERING: hold until we know it is not a tool call. + content_buffer += delta + stripped = content_buffer.lstrip() + if not stripped: + continue + + is_match = False + is_prefix = False + for sig in TOOL_XML_SIGNALS: + if stripped.startswith(sig): + is_match = True + break + if sig.startswith(stripped): + is_prefix = True + break + + if is_match: + detect_state = _state_draining + elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS: + continue + else: + detect_state = _state_streaming + cumulative_display += content_buffer + cleaned = strip_tool_markup(cumulative_display) + if len(cleaned) > len(last_emitted): + last_emitted = cleaned + yield {"type": "content", "text": cleaned} + + # Stream finished -- resolve what we collected. + if cancel_event is not None and cancel_event.is_set(): + return + + if detect_state == _state_buffering: + # Buffer never resolved -- tool XML or plain content. + stripped = content_buffer.lstrip() + if stripped and has_tool_signal(stripped): + detect_state = _state_draining + else: + if content_buffer: + cumulative_display += content_buffer + yield { + "type": "content", + "text": strip_tool_markup(cumulative_display, final = True), + } + yield {"type": "status", "text": ""} + return + + if detect_state == _state_streaming: + # No tool detected mid-stream -- check for late tool XML. + safety_tc = None + if has_tool_signal(content_accum): + safety_tc = parse_tool_calls_from_text( + content_accum, + id_offset = next_call_id, + ) + if not safety_tc: + # Final answer: streaming already emitted content. + # Skip a final=True re-strip so literal "" + # in prose survives when no real tool call parsed. + yield {"type": "status", "text": ""} + return + tool_calls = safety_tc + content_text = strip_tool_markup(content_accum, final = True) + logger.info( + "Safetensors safety net: parsed %d tool call(s) from streamed content", + len(tool_calls), + ) + else: + # DRAINING: parse tool calls out of full content. + tool_calls = parse_tool_calls_from_text( + content_accum, + id_offset = next_call_id, + ) + if not tool_calls and auto_heal_tool_calls: + # Parser found nothing -- surface raw content so any + # literal "" prose is preserved. + if content_accum: + yield {"type": "content", "text": content_accum} + yield {"type": "status", "text": ""} + return + content_text = strip_tool_markup(content_accum, final = True) + + if final_attempt_done: + # Final-answer turn re-called a tool -- stop the loop. + if content_text: + yield {"type": "content", "text": content_text} + yield {"type": "status", "text": ""} + return + + assistant_msg: dict = {"role": "assistant", "content": content_text} + if tool_calls: + assistant_msg["tool_calls"] = tool_calls + next_call_id += len(tool_calls) + conversation.append(assistant_msg) + + for tc in tool_calls or []: + func = tc.get("function", {}) or {} + tool_name = func.get("name", "") or "" + arguments = _coerce_arguments( + func.get("arguments", {}), + heal = auto_heal_tool_calls, + tool_name = tool_name, + ) + + yield {"type": "status", "text": _status_for_tool(tool_name, arguments)} + yield { + "type": "tool_start", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "arguments": arguments, + } + + tc_key = tool_name + str(arguments) + if allowed_tool_names and tool_name not in allowed_tool_names: + result = ( + f"Error: tool '{tool_name}' is not enabled for this " + "request. Use one of the enabled tools or provide a " + "final answer." + ) + else: + already_ran_ok = any( + k == tc_key and not err for k, err in tool_call_history + ) + if already_ran_ok: + result = DUPLICATE_CALL_NUDGE + else: + eff_timeout = ( + None if tool_call_timeout >= 9999 else tool_call_timeout + ) + try: + result = execute_tool( + tool_name, + arguments, + cancel_event = cancel_event, + timeout = eff_timeout, + session_id = session_id, + ) + except Exception as exc: + logger.exception("Tool %s raised: %s", tool_name, exc) + result = f"Error: tool raised an exception: {exc}" + + yield { + "type": "tool_end", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "result": result, + } + + is_error = isinstance(result, str) and result.lstrip().startswith( + TOOL_ERROR_PREFIXES + ) + tool_call_history.append((tc_key, is_error)) + + # Strip frontend image sentinel from the model's view. + # Cut at the first occurrence so leading and consecutive + # sentinels are both removed. + result_for_model = result + if isinstance(result_for_model, str) and "__IMAGES__:" in result_for_model: + result_for_model = result_for_model.split("__IMAGES__:", 1)[0].rstrip() + if is_error: + result_for_model = result_for_model + TOOL_ERROR_NUDGE + + tool_msg: dict = { + "role": "tool", + "name": tool_name, + "content": result_for_model, + } + tool_call_id = tc.get("id") + if tool_call_id: + tool_msg["tool_call_id"] = tool_call_id + conversation.append(tool_msg) + + # Clear the status badge before the next turn. + yield {"type": "status", "text": ""} + + if iteration + 1 >= max_tool_iterations and not final_attempt_done: + # Budget exhausted; nudge a final plain answer. + final_attempt_done = True + conversation.append( + { + "role": "user", + "content": BUDGET_EXHAUSTED_NUDGE, + } + ) + + yield {"type": "status", "text": ""} diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py new file mode 100644 index 0000000000..a0ab8a2a53 --- /dev/null +++ b/studio/backend/core/inference/tool_call_parser.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Backend-neutral tool-call XML parser shared by GGUF and safetensors. +Tolerates missing closing tags in either ``{json}`` +or ``v...`` shape. +""" + +import json +import re + + +# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing +# unclosed runs so truncated tails don't leak markup. +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + re.compile(r".*?", re.DOTALL), +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r".*$", re.DOTALL), + re.compile(r".*$", re.DOTALL), +] + + +# Prefixes the streaming buffer watches for to gate in-progress text. +TOOL_XML_SIGNALS = ("", "\s*\{") +_TC_FUNC_START_RE = re.compile(r"\s*") +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +_TC_PARAM_START_RE = re.compile(r"\s*") +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + + +def strip_tool_markup(text: str, *, final: bool = False) -> str: + """Strip tool-call XML from streamed text. + + ``final=False`` only removes closed pairs (used during streaming so + in-progress XML stays buffered). ``final=True`` also removes a + trailing unclosed run and trims the result. + """ + pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + for pat in pats: + text = pat.sub("", text) + return text.strip() if final else text + + +def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict]: + """Parse OpenAI-format ``tool_calls`` from model text. + + Returns a list of ``{"id", "type", "function": {"name", "arguments"}}`` + dicts. ``arguments`` is always a JSON string so callers can hand it + straight back into an OpenAI-style response. + + Handles two shapes: + + - JSON inside ```` tags: + ``{"name":"web_search","arguments":{"query":"..."}}`` + - XML-style function blocks: + ``v`` + + Closing tags (````, ````, ````) + are all optional since models frequently omit them. + """ + tool_calls: list[dict] = [] + + # Pattern 1: {json}. Balanced-brace scan that skips + # braces inside JSON strings. + for m in _TC_JSON_START_RE.finditer(content): + brace_start = m.end() - 1 # position of the opening { + depth, i = 0, brace_start + in_string = False + while i < len(content): + ch = content[i] + if in_string: + if ch == "\\" and i + 1 < len(content): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + break + i += 1 + if depth == 0: + json_str = content[brace_start : i + 1] + try: + obj = json.loads(json_str) + tc = { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": obj.get("name", ""), + "arguments": obj.get("arguments", {}), + }, + } + if isinstance(tc["function"]["arguments"], dict): + tc["function"]["arguments"] = json.dumps( + tc["function"]["arguments"] + ) + tool_calls.append(tc) + except (json.JSONDecodeError, ValueError): + pass + + # Pattern 2: v... -- closing tags + # optional; don't use as body boundary because code + # values can contain that literal. + if not tool_calls: + func_starts = list(_TC_FUNC_START_RE.finditer(content)) + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + next_func = ( + func_starts[idx + 1].start() + if idx + 1 < len(func_starts) + else len(content) + ) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = content[body_start:body_end] + body = _TC_FUNC_CLOSE_RE.sub("", body) + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + # Single param: take everything to body end so + # embedded in code strings is preserved. + pm = param_starts[0] + val = body[pm.end() :] + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = val.strip() + else: + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() + if pidx + 1 < len(param_starts) + else len(body) + ) + val = body[val_start:next_param] + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = val.strip() + + tc = { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": func_name, + "arguments": json.dumps(arguments), + }, + } + tool_calls.append(tc) + + return tool_calls + + +def has_tool_signal(text: str) -> bool: + """Return True if ``text`` contains any tool-call XML signal.""" + return any(s in text for s in TOOL_XML_SIGNALS) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index cacede2d3e..20a7d2d16c 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -346,6 +346,26 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: "audio_type": getattr(mc, "audio_type", None), "has_audio_input": getattr(mc, "has_audio_input", False), } + # Forward chat_template_info so the parent can classify + # capabilities without re-entering the subprocess. + try: + _bm = getattr(backend, "models", {}) or {} + _entry = ( + _bm.get(mc.identifier) + or _bm.get(getattr(backend, "active_model_name", None)) + or {} + ) + _tpl_info = _entry.get("chat_template_info") + if isinstance(_tpl_info, dict): + model_info["chat_template_info"] = { + "has_template": bool(_tpl_info.get("has_template", False)), + "template": _tpl_info.get("template"), + "format_type": _tpl_info.get("format_type", "generic"), + "template_name": _tpl_info.get("template_name"), + "special_tokens": _tpl_info.get("special_tokens", {}) or {}, + } + except Exception as _tpl_exc: + logger.warning("chat_template_info forward failed: %s", _tpl_exc) _send_response( resp_queue, { @@ -416,6 +436,18 @@ def _handle_generate( "cancel_event": cancel_event, } + # Optional template/tool plumbing: only forward keys that are + # actually present so the backend signature can evolve without + # breaking older command payloads. + for opt_key in ( + "tools", + "enable_thinking", + "reasoning_effort", + "preserve_thinking", + ): + if opt_key in cmd: + gen_kwargs[opt_key] = cmd[opt_key] + # Choose generation path use_adapter = cmd.get("use_adapter") if use_adapter is not None: @@ -648,36 +680,6 @@ def run_inference_process( os.environ["HF_HUB_DISABLE_XET"] = "1" logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)") - # Offline auto-detect: skip 25s of hf_hub_download retries per file - # if DNS is dead; cached files resolve instantly under HF_HUB_OFFLINE=1. - # Scope is this subprocess only -- orchestrator spawns a fresh worker - # per load (see core/inference/orchestrator.py), so the env cannot - # persist across loads. - if "HF_HUB_OFFLINE" not in os.environ: - import socket as _socket - import threading as _threading - - # Probe on a daemon thread so concurrent sockets in the parent - # interpreter are not affected by socket.setdefaulttimeout. - _result: list = [None] - - def _probe() -> None: - try: - _socket.gethostbyname("huggingface.co") - _result[0] = False - except Exception: - _result[0] = True - - _t = _threading.Thread(target = _probe, daemon = True) - _t.start() - _t.join(2.0) - if _result[0] is None or _result[0] is True: - os.environ["HF_HUB_OFFLINE"] = "1" - os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") - logger.warning( - "huggingface.co unreachable; HF_HUB_OFFLINE=1 set for this worker." - ) - import warnings from loggers.config import LogConfig diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2f095cb823..1b4e7051b0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -235,6 +235,57 @@ router = APIRouter() studio_router = APIRouter() +def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: + """Classify reasoning/tool capabilities via the GGUF classifier so + flags match across backends. gpt-oss is overridden because Harmony + routes reasoning and tools through tokenizer channels, not template + markup.""" + model_id = getattr(backend, "active_model_name", None) + flags = ( + detect_reasoning_flags( + chat_template, + model_identifier = model_id, + log_source = "safetensors", + ) + if chat_template + else { + "supports_reasoning": False, + "reasoning_style": "enable_thinking", + "reasoning_always_on": False, + "supports_preserve_thinking": False, + "supports_tools": False, + } + ) + # Our safetensors loop only parses {json} + # and .... Llama uses <|python_tag|>, + # Mistral uses [TOOL_CALLS]; advertising tools for those would + # enable a pill the parser cannot honour. GGUF is unaffected -- + # llama-server normalises every format into structured deltas. + if ( + flags.get("supports_tools") + and chat_template + and "" not in chat_template + and " XML this loop parses). + try: + if hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model(): + flags["supports_reasoning"] = True + flags["reasoning_style"] = "reasoning_effort" + flags["supports_tools"] = False + except Exception: + logger.debug("gpt_oss_check_failed", exc_info = True) + return flags + + def _effective_enable_tools(payload) -> Optional[bool]: """Resolve `payload.enable_tools` against the process-level tool policy. @@ -590,6 +641,7 @@ async def load_model( reasoning_style = llama_backend.reasoning_style, reasoning_always_on = llama_backend.reasoning_always_on, supports_preserve_thinking = llama_backend.supports_preserve_thinking, + supports_tools = llama_backend.supports_tools, chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, @@ -612,21 +664,10 @@ async def load_model( logger.warning( f"Could not retrieve chat template for {backend.active_model_name}: {e}" ) - # Non-GGUF: only advertise reasoning for gpt-oss Harmony, - # which emits reasoning via channels at the tokenizer level. - # Template-level chat_template_kwargs (enable_thinking / - # preserve_thinking / tools) are not yet forwarded through - # the transformers generation path, so avoid advertising - # controls the server cannot honour outside GGUF. - _sf_supports_reasoning = False - _sf_reasoning_style = "enable_thinking" - if hasattr(backend, "_is_gpt_oss_model"): - try: - if backend._is_gpt_oss_model(): - _sf_supports_reasoning = True - _sf_reasoning_style = "reasoning_effort" - except Exception: - pass + # Classify via the same path as GGUF. + _sf_flags = _detect_safetensors_features(backend, _chat_template) + _sf_supports_reasoning = _sf_flags["supports_reasoning"] + _sf_reasoning_style = _sf_flags["reasoning_style"] return LoadResponse( status = "already_loaded", model = model_log_label @@ -647,9 +688,9 @@ async def load_model( ), supports_reasoning = _sf_supports_reasoning, reasoning_style = _sf_reasoning_style, - reasoning_always_on = False, - supports_preserve_thinking = False, - supports_tools = False, + reasoning_always_on = _sf_flags["reasoning_always_on"], + supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], + supports_tools = _sf_flags["supports_tools"], chat_template = _chat_template, ) @@ -982,19 +1023,8 @@ async def load_model( except Exception: pass - # Non-GGUF: gpt-oss Harmony surfaces reasoning via tokenizer-level - # channels; other safetensors reasoning/tools/preserve-thinking - # knobs are not forwarded to tokenizer.apply_chat_template yet, so - # we only advertise support for the Harmony case here. - _sf_supports_reasoning = False - _sf_reasoning_style = "enable_thinking" - if hasattr(backend, "_is_gpt_oss_model"): - try: - if backend._is_gpt_oss_model(): - _sf_supports_reasoning = True - _sf_reasoning_style = "reasoning_effort" - except Exception: - pass + # Classify reasoning/tool flags via the GGUF sniffer. + _sf_flags = _detect_safetensors_features(backend, _chat_template) return LoadResponse( status = "loaded", @@ -1012,11 +1042,11 @@ async def load_model( requires_trust_remote_code = bool( inference_config.get("trust_remote_code", False) ), - supports_reasoning = _sf_supports_reasoning, - reasoning_style = _sf_reasoning_style, - reasoning_always_on = False, - supports_preserve_thinking = False, - supports_tools = False, + supports_reasoning = _sf_flags["supports_reasoning"], + reasoning_style = _sf_flags["reasoning_style"], + reasoning_always_on = _sf_flags["reasoning_always_on"], + supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], + supports_tools = _sf_flags["supports_tools"], chat_template = _chat_template, ) @@ -1388,18 +1418,8 @@ async def get_status( else None ) - # Non-GGUF: only gpt-oss Harmony is wired through the transformers - # generation path. Other template-level reasoning / tool kwargs - # are not yet forwarded, so we do not advertise them here. - supports_reasoning = False - reasoning_style = "enable_thinking" - if backend.active_model_name and hasattr(backend, "_is_gpt_oss_model"): - try: - if backend._is_gpt_oss_model(): - supports_reasoning = True - reasoning_style = "reasoning_effort" - except Exception: - pass + # Non-GGUF: classify from the loaded template. + _sf_flags = _detect_safetensors_features(backend, chat_template) inference_config = ( load_inference_config(backend.active_model_name) if backend.active_model_name @@ -1419,11 +1439,11 @@ async def get_status( requires_trust_remote_code = bool( (inference_config or {}).get("trust_remote_code", False) ), - supports_reasoning = supports_reasoning, - reasoning_style = reasoning_style, - reasoning_always_on = False, - supports_preserve_thinking = False, - supports_tools = False, + supports_reasoning = _sf_flags["supports_reasoning"], + reasoning_style = _sf_flags["reasoning_style"], + reasoning_always_on = _sf_flags["reasoning_always_on"], + supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], + supports_tools = _sf_flags["supports_tools"], chat_template = chat_template, llama_cpp_supports_mtp = _supports_mtp, llama_cpp_prebuilt_stale = _stale, @@ -2749,6 +2769,300 @@ async def openai_chat_completions( except Exception as e: raise HTTPException(status_code = 400, detail = f"Failed to decode image: {e}") + # Classify capability flags from the loaded template. + _sf_model_info = backend.models.get(backend.active_model_name, {}) + _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") + _sf_features = _detect_safetensors_features(backend, _sf_tpl) + + cancel_event = threading.Event() + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + created = int(time.time()) + + # ── Safetensors tool-calling path ───────────────────────── + # Mirrors the GGUF agentic loop's event shape. Disabled for + # vision turns (untested overlap with image render slot) and + # for gpt-oss (Harmony uses dedicated channels, not + # XML -- gpt-oss tools still work via the GGUF path). + _sf_is_gptoss = False + try: + _sf_is_gptoss = bool( + hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model() + ) + except Exception: + _sf_is_gptoss = False + + _sf_tool_budget = ( + payload.max_tool_calls_per_message + if payload.max_tool_calls_per_message is not None + else 25 + ) + + _sf_use_tools = ( + _effective_enable_tools(payload) + and _sf_features.get("supports_tools", False) + and image is None + and not _sf_is_gptoss + and _sf_tool_budget > 0 + ) + + if _sf_use_tools: + from core.inference.tools import ALL_TOOLS + + if payload.enabled_tools is not None: + _sf_tools_to_use = [ + t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools + ] + else: + _sf_tools_to_use = ALL_TOOLS + + _sf_tool_names = {t["function"]["name"] for t in _sf_tools_to_use} + _sf_has_web = "web_search" in _sf_tool_names + _sf_has_code = "python" in _sf_tool_names or "terminal" in _sf_tool_names + + _sf_date_line = f"The current date is {_date.today().isoformat()}." + _sf_model_size_b = _extract_model_size_b(model_name) + _sf_is_small_model = _sf_model_size_b is not None and _sf_model_size_b < 9 + + if _sf_is_small_model: + _sf_web_tips = "Do not repeat the same search query." + else: + _sf_web_tips = ( + "When you search and find a relevant URL in the results, " + "fetch its full content by calling web_search with the url parameter. " + "Do not repeat the same search query. If a search returns " + "no useful results, try rephrasing or fetching a result URL directly." + ) + _sf_code_tips = ( + "Use code execution for math, calculations, data processing, " + "or to parse and analyze information from tool results." + ) + + if _sf_has_web and _sf_has_code: + _sf_nudge = ( + _sf_date_line + " " + "You have access to tools. When appropriate, prefer using " + "tools rather than answering from memory. " + + _sf_web_tips + + " " + + _sf_code_tips + ) + elif _sf_has_code: + _sf_nudge = ( + _sf_date_line + " " + "You have access to tools. When appropriate, prefer using " + "code execution rather than answering from memory. " + _sf_code_tips + ) + elif _sf_has_web: + _sf_nudge = ( + _sf_date_line + " " + "You have access to tools. When appropriate, prefer using " + "web search for up-to-date or uncertain factual " + "information rather than answering from memory. " + _sf_web_tips + ) + else: + _sf_nudge = "" + + _sf_system_prompt = system_prompt + if _sf_nudge: + _sf_nudge += _TOOL_ACTION_NUDGE + if _sf_system_prompt: + _sf_system_prompt = _sf_system_prompt.rstrip() + "\n\n" + _sf_nudge + else: + _sf_system_prompt = _sf_nudge + + # Strip stale tool-call XML from prior assistant turns. + _sf_chat_messages = [] + for _msg in chat_messages: + if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + _sf_chat_messages.append( + { + **_msg, + "content": _TOOL_XML_RE.sub("", _msg["content"]).strip(), + } + ) + else: + _sf_chat_messages.append(_msg) + + def sf_generate_with_tools(): + return backend.generate_chat_completion_with_tools( + messages = _sf_chat_messages, + tools = _sf_tools_to_use, + system_prompt = _sf_system_prompt or "", + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_tokens = payload.max_tokens, + repetition_penalty = payload.repetition_penalty, + cancel_event = cancel_event, + enable_thinking = payload.enable_thinking, + reasoning_effort = payload.reasoning_effort, + preserve_thinking = payload.preserve_thinking, + auto_heal_tool_calls = payload.auto_heal_tool_calls + if payload.auto_heal_tool_calls is not None + else True, + max_tool_iterations = _sf_tool_budget, + tool_call_timeout = payload.tool_call_timeout + if payload.tool_call_timeout is not None + else 300, + session_id = payload.session_id, + use_adapter = payload.use_adapter, + ) + + _sf_tool_sentinel = object() + _sf_cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _sf_tracker = _TrackedCancel(cancel_event, *_sf_cancel_keys) + _sf_tracker.__enter__() + + async def sf_tool_stream(): + try: + first_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + ], + ) + yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" + + gen = sf_generate_with_tools() + prev_text = "" + while True: + if cancel_event.is_set(): + backend.reset_generation_state() + break + if await request.is_disconnected(): + cancel_event.set() + backend.reset_generation_state() + return + + event = await asyncio.to_thread(next, gen, _sf_tool_sentinel) + if event is _sf_tool_sentinel: + break + + if event["type"] == "status": + if not event["text"]: + prev_text = "" + status_data = json.dumps( + { + "type": "tool_status", + "content": event["text"], + } + ) + yield f"data: {status_data}\n\n" + continue + + if event["type"] in ("tool_start", "tool_end"): + if event["type"] == "tool_start": + prev_text = "" + yield f"data: {json.dumps(event)}\n\n" + continue + + # Diff cumulative cleaned text against last snapshot. + raw_cumulative = event.get("text", "") + clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative) + new_text = clean_cumulative[len(prev_text) :] + prev_text = clean_cumulative + if not new_text: + continue + chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(content = new_text), + finish_reason = None, + ) + ], + ) + yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + + final_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = "stop", + ) + ], + ) + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + yield "data: [DONE]\n\n" + + except asyncio.CancelledError: + cancel_event.set() + backend.reset_generation_state() + raise + except Exception: + backend.reset_generation_state() + # Generic wire message; full trace stays in the log + # (CWE-209: transformers/torch errors may leak paths). + logger.exception("safetensors tool stream error") + error_chunk = { + "error": { + "message": "An internal error occurred.", + "type": "server_error", + }, + } + yield f"data: {json.dumps(error_chunk)}\n\n" + finally: + _sf_tracker.__exit__(None, None, None) + + if payload.stream: + return StreamingResponse( + sf_tool_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + # Non-streaming JSON: drain the loop, build one ChatCompletion. + try: + + def _drain_to_text(): + full_text = "" + gen = sf_generate_with_tools() + for event in gen: + if cancel_event.is_set(): + break + if event.get("type") == "content": + full_text = _TOOL_XML_RE.sub("", event.get("text", "")) + return full_text + + content_text = await asyncio.to_thread(_drain_to_text) + response = ChatCompletion( + id = completion_id, + created = created, + model = model_name, + choices = [ + CompletionChoice( + message = CompletionMessage(content = content_text), + finish_reason = "stop", + ) + ], + ) + return JSONResponse(content = response.model_dump()) + except Exception: + backend.reset_generation_state() + # CWE-209: generic detail; full trace in log. + logger.exception("safetensors tool completion error") + raise HTTPException( + status_code = 500, + detail = "An internal error occurred.", + ) + finally: + _sf_tracker.__exit__(None, None, None) + # Shared generation kwargs gen_kwargs = dict( messages = chat_messages, @@ -2761,9 +3075,14 @@ async def openai_chat_completions( max_new_tokens = payload.max_tokens or 2048, repetition_penalty = payload.repetition_penalty, ) - - # Choose generation path (adapter-controlled or standard) - cancel_event = threading.Event() + # Forward reasoning kwargs; the worker/template wrapper peels off + # any the template doesn't accept. + if payload.enable_thinking is not None: + gen_kwargs["enable_thinking"] = payload.enable_thinking + if payload.reasoning_effort is not None: + gen_kwargs["reasoning_effort"] = payload.reasoning_effort + if payload.preserve_thinking is not None: + gen_kwargs["preserve_thinking"] = payload.preserve_thinking if payload.use_adapter is not None: @@ -2780,9 +3099,6 @@ async def openai_chat_completions( cancel_event = cancel_event, **gen_kwargs ) - completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" - created = int(time.time()) - # ── Streaming response ──────────────────────────────────────── if payload.stream: _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index ce447bdd1f..16cca7dd40 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -158,3 +158,97 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert backend._is_vlm is True assert isinstance(backend._processor, _DummyProcessor) assert isinstance(backend._tokenizer, _DummyTokenizer) + + +# Regression: MLXInferenceBackend.generate_chat_response must accept the +# four template kwargs (tools / enable_thinking / reasoning_effort / +# preserve_thinking) so the route layer can forward what the user +# toggled in the UI. The previous signature raised +# "got an unexpected keyword argument 'tools'" on Mac. + + +def test_mlx_generate_chat_response_accepts_template_kwargs(): + import inspect + from core.inference.mlx_inference import MLXInferenceBackend + + sig = inspect.signature(MLXInferenceBackend.generate_chat_response) + params = sig.parameters + for name in ("tools", "enable_thinking", "reasoning_effort", "preserve_thinking"): + assert name in params, ( + f"MLX.generate_chat_response is missing the {name!r} kwarg; " + "the route layer forwards this and a missing kwarg raises " + "TypeError on Mac" + ) + assert ( + params[name].default is None + ), f"{name!r} must default to None so existing callers stay valid" + + +def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): + """The Mac text path must route through apply_chat_template_for_ + generation so reasoning / tool kwargs reach the tokenizer.""" + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + captured = {} + + def _fake_apply(tokenizer, messages, **kwargs): + captured["tokenizer"] = tokenizer + captured["messages"] = messages + captured["kwargs"] = kwargs + return "" + + monkeypatch.setattr( + "core.inference.chat_template_helpers." "apply_chat_template_for_generation", + _fake_apply, + raising = True, + ) + + # mlx_lm.stream_generate yields response objects with .token; make a + # one-token generator so _generate_text returns without touching the + # real stack. + import types as _types + + mlx_lm_pkg = _types.ModuleType("mlx_lm") + mlx_lm_sample = _types.ModuleType("mlx_lm.sample_utils") + mlx_lm_sample.make_sampler = lambda **_kw: object() + mlx_lm_sample.make_logits_processors = lambda **_kw: None + + class _Resp: + def __init__(self, tok): + self.token = tok + + def _stream_generate(_model, _tokenizer, **_kw): + yield _Resp(1) + + mlx_lm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample) + + class _Tok: + chat_template = "x" + + def decode(self, ids, skip_special_tokens = False): + return "hi" + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = _Tok() + backend._is_vlm = False + + out = list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + tools = [{"function": {"name": "web_search"}}], + enable_thinking = True, + reasoning_effort = "medium", + preserve_thinking = True, + max_new_tokens = 1, + ) + ) + assert out == ["hi"] + # The kwargs the user toggled must reach the chat-template helper. + assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}] + assert captured["kwargs"]["enable_thinking"] is True + assert captured["kwargs"]["reasoning_effort"] == "medium" + assert captured["kwargs"]["preserve_thinking"] is True diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py new file mode 100644 index 0000000000..c3ee5b9ff1 --- /dev/null +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -0,0 +1,451 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Capability advertisement contract: classifier honesty, worker→ +orchestrator IPC hop, and route-layer end-to-end. Pure helpers + fakes; +no torch / transformers import. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +# Qwen3 snippet covering tools, enable_thinking, preserve_thinking. +QWEN3_TEMPLATE = """ +{%- if tools %} + {{- '<|im_start|>system\\nFor each function call, return a json object' + ' wrapped inside tags.\\n' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'tool' %} + {{- '<|im_start|>tool\\n' + message.content + '<|im_end|>\\n' }} + {%- endif %} +{%- endfor %} +{%- if enable_thinking is defined and enable_thinking %} + {{- '' }} +{%- endif %} +{%- if preserve_thinking %} + {{- assistant.reasoning_content }} +{%- endif %} +""" + + +GPT_OSS_TEMPLATE = """ +<|start|>system<|message|>You are gpt-oss. +reasoning_effort: {{ reasoning_effort }} +<|end|> +""" + + +PLAIN_TEMPLATE = """ +{%- for message in messages %} + {{- message.role + ': ' + message.content + '\\n' }} +{%- endfor %} +""" + + +# ── Tests: classifier honesty ──────────────────────────────────────── + + +def test_detect_reasoning_flags_qwen3_supports_tools_and_reasoning(): + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(QWEN3_TEMPLATE, "unsloth/Qwen3-0.6B") + assert flags["supports_tools"] is True + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking" + assert flags["supports_preserve_thinking"] is True + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_plain_template_all_false(): + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(PLAIN_TEMPLATE, "some/PlainChat") + assert flags["supports_tools"] is False + assert flags["supports_reasoning"] is False + assert flags["supports_preserve_thinking"] is False + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_none_template_returns_all_false(): + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(None) + assert flags["supports_tools"] is False + assert flags["supports_reasoning"] is False + assert flags["supports_preserve_thinking"] is False + assert flags["reasoning_always_on"] is False + assert flags["reasoning_style"] == "enable_thinking" + + +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 + + backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") + flags = _detect_safetensors_features(backend, QWEN3_TEMPLATE) + assert flags["supports_tools"] is True + assert flags["supports_reasoning"] is True + + +def test_detect_safetensors_features_none_template_returns_all_false(): + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") + flags = _detect_safetensors_features(backend, None) + assert flags == { + "supports_reasoning": False, + "reasoning_style": "enable_thinking", + "reasoning_always_on": False, + "supports_preserve_thinking": False, + "supports_tools": False, + } + + +def test_detect_safetensors_features_gptoss_disables_tools(): + """gpt-oss Harmony: tools intentionally off even if template marks it.""" + from routes.inference import _detect_safetensors_features + + backend = MagicMock() + backend.active_model_name = "unsloth/gpt-oss-20b" + backend._is_gpt_oss_model.return_value = True + + flags = _detect_safetensors_features(backend, QWEN3_TEMPLATE) + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "reasoning_effort" + assert flags["supports_tools"] is False + + +# Llama-3 / Mistral templates advertise tool handling but the model emits +# tool calls in <|python_tag|> / [TOOL_CALLS] format -- not the +# / system<|end_header_id|>' }} + {{- 'You have access to the following tools.' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'tool' %} + {{- '<|start_header_id|>ipython<|end_header_id|>' }} + {{- '<|python_tag|>' }} + {{- message.content }} + {%- endif %} +{%- endfor %} +""" + +MISTRAL_TEMPLATE = """ +{%- if tools %} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'tool' %} + {{- '[TOOL_CALLS]' + message.content + '[/TOOL_CALLS]' }} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_llama3_template_suppresses_tools(): + """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE) + assert flags["supports_tools"] is False + + +def test_detect_safetensors_features_mistral_template_suppresses_tools(): + """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") + flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE) + assert flags["supports_tools"] is False + + +def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on(): + """Sanity check: gate only suppresses non-Qwen formats.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") + flags = _detect_safetensors_features(backend, QWEN3_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_function_xml_format_keeps_tools_on(): + """Templates emitting XML are parser-compatible.""" + from routes.inference import _detect_safetensors_features + + tpl_with_function_xml = ( + "{%- if tools %}<|im_start|>system\n" + "Tool call format: v" + "<|im_end|>{%- endif %}" + ) + backend = SimpleNamespace(active_model_name = "custom/with-function-xml") + flags = _detect_safetensors_features(backend, tpl_with_function_xml) + assert flags["supports_tools"] is True + + +# Qwen3.5 family pins -- the live GGUF + safetensors templates fetched +# from the unsloth/Qwen3.5-0.8B(-GGUF) repos both wrap tool calls as +# ``\n...``. Capture a faithful slice so the +# classifier never silently regresses for this family. + +QWEN35_TOOL_INSTRUCTION = ( + "{%- if tools %}\n" + " <|im_start|>system\n" + " # Tools\n" + " \n" + " {%- for tool in tools %}{{ tool | tojson }}{%- endfor %}\n" + " \n" + " If you choose to call a function ONLY reply in the following format:\n" + " \n" + " \n" + " \n" + " value_1\n" + " \n" + " \n" + " \n" + " <|im_end|>\n" + "{%- endif %}\n" + "{%- if enable_thinking is defined and enable_thinking %}{{- '' }}{%- endif %}\n" +) + + +def test_detect_safetensors_features_qwen35_keeps_tools_on(): + """unsloth/Qwen3.5-0.8B family must surface tools+reasoning enabled.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Qwen3.5-0.8B") + flags = _detect_safetensors_features(backend, QWEN35_TOOL_INSTRUCTION) + assert flags["supports_tools"] is True + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking" + + +# ── Tests: IPC bridge contract ─────────────────────────────────────── + + +def test_orchestrator_mirrors_chat_template_info_into_models_dict(): + """Worker → orchestrator must copy chat_template_info verbatim.""" + from core.inference.orchestrator import InferenceOrchestrator + + orch = InferenceOrchestrator.__new__(InferenceOrchestrator) + orch.models = {} + orch.active_model_name = None + orch.loading_models = set() + + model_info = { + "identifier": "unsloth/Qwen3-0.6B", + "display_name": "Qwen3-0.6B", + "is_vision": False, + "is_lora": False, + "is_gguf": False, + "is_audio": False, + "audio_type": None, + "has_audio_input": False, + "chat_template_info": { + "has_template": True, + "template": QWEN3_TEMPLATE, + "format_type": "chatml", + "template_name": "qwen3", + "special_tokens": {"bos_token": "<|im_start|>"}, + }, + } + + # Replay orchestrator.load_model's mirror block verbatim. + orch.active_model_name = model_info["identifier"] + orch.models[orch.active_model_name] = { + "is_vision": model_info.get("is_vision", False), + "is_lora": model_info.get("is_lora", False), + "display_name": model_info.get("display_name", "x"), + "is_audio": model_info.get("is_audio", False), + "audio_type": model_info.get("audio_type"), + "has_audio_input": model_info.get("has_audio_input", False), + } + _tpl_info = model_info.get("chat_template_info") + if isinstance(_tpl_info, dict): + orch.models[orch.active_model_name]["chat_template_info"] = _tpl_info + + entry = orch.models[orch.active_model_name] + tpl = entry.get("chat_template_info", {}).get("template") + assert tpl == QWEN3_TEMPLATE + + from routes.inference import _detect_safetensors_features + + flags = _detect_safetensors_features( + SimpleNamespace(active_model_name = orch.active_model_name), tpl + ) + assert flags["supports_tools"] is True + assert flags["supports_reasoning"] is True + + +def test_orchestrator_missing_chat_template_info_falls_back_to_all_false(): + """Old / malformed worker reply: no crash, all flags False.""" + from core.inference.orchestrator import InferenceOrchestrator + from routes.inference import _detect_safetensors_features + + orch = InferenceOrchestrator.__new__(InferenceOrchestrator) + orch.models = {} + orch.active_model_name = "unsloth/Qwen3-0.6B" + + model_info = { + "identifier": "unsloth/Qwen3-0.6B", + "is_vision": False, + "is_lora": False, + # NB: no chat_template_info key + } + orch.models[orch.active_model_name] = { + "is_vision": False, + "is_lora": False, + } + _tpl_info = model_info.get("chat_template_info") + if isinstance(_tpl_info, dict): + orch.models[orch.active_model_name]["chat_template_info"] = _tpl_info + + entry = orch.models[orch.active_model_name] + tpl = entry.get("chat_template_info", {}).get("template") + assert tpl is None + + flags = _detect_safetensors_features( + SimpleNamespace(active_model_name = orch.active_model_name), tpl + ) + assert flags["supports_tools"] is False + + +def test_worker_load_reply_payload_includes_chat_template_info(): + """Worker IPC reply carries chat_template_info dict.""" + + class _StubBackend: + def __init__(self, identifier, template): + self.active_model_name = identifier + self.models = { + identifier: { + "chat_template_info": { + "has_template": True, + "template": template, + "format_type": "chatml", + "template_name": "qwen3", + "special_tokens": {"bos_token": "<|im_start|>"}, + } + } + } + + backend = _StubBackend("unsloth/Qwen3-0.6B", QWEN3_TEMPLATE) + mc = SimpleNamespace( + identifier = "unsloth/Qwen3-0.6B", + display_name = "Qwen3-0.6B", + is_vision = False, + is_lora = False, + ) + + # Replay the worker's payload-build block. + model_info = { + "identifier": mc.identifier, + "display_name": mc.display_name, + "is_vision": mc.is_vision, + "is_lora": mc.is_lora, + "is_gguf": False, + } + _bm = getattr(backend, "models", {}) or {} + _entry = ( + _bm.get(mc.identifier) + or _bm.get(getattr(backend, "active_model_name", None)) + or {} + ) + _tpl_info = _entry.get("chat_template_info") + if isinstance(_tpl_info, dict): + model_info["chat_template_info"] = { + "has_template": bool(_tpl_info.get("has_template", False)), + "template": _tpl_info.get("template"), + "format_type": _tpl_info.get("format_type", "generic"), + "template_name": _tpl_info.get("template_name"), + "special_tokens": _tpl_info.get("special_tokens", {}) or {}, + } + + assert "chat_template_info" in model_info + assert model_info["chat_template_info"]["template"] == QWEN3_TEMPLATE + assert model_info["chat_template_info"]["has_template"] is True + + +def test_worker_load_reply_payload_survives_missing_template(): + """Tokenizer with no chat_template still produces a valid reply.""" + + class _StubBackend: + def __init__(self): + self.active_model_name = "legacy/no-template" + self.models = {"legacy/no-template": {}} # no chat_template_info + + backend = _StubBackend() + mc = SimpleNamespace( + identifier = "legacy/no-template", + display_name = "legacy", + is_vision = False, + is_lora = False, + ) + + model_info = { + "identifier": mc.identifier, + "display_name": mc.display_name, + "is_vision": mc.is_vision, + "is_lora": mc.is_lora, + "is_gguf": False, + } + _bm = getattr(backend, "models", {}) or {} + _entry = _bm.get(mc.identifier) or {} + _tpl_info = _entry.get("chat_template_info") + if isinstance(_tpl_info, dict): + model_info["chat_template_info"] = dict(_tpl_info) + + assert "chat_template_info" not in model_info + + +# ── End-to-end: route layer sees the template, advertises True ─────── + + +def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): + """End-to-end: Qwen3 safetensors flips supports_tools=True.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace( + active_model_name = "unsloth/Qwen3-0.6B", + models = { + "unsloth/Qwen3-0.6B": { + "is_vision": False, + "chat_template_info": { + "has_template": True, + "template": QWEN3_TEMPLATE, + "format_type": "chatml", + }, + } + }, + ) + + _model_info = backend.models.get(backend.active_model_name, {}) + _tpl = _model_info.get("chat_template_info", {}).get("template") + flags = _detect_safetensors_features(backend, _tpl) + + assert flags["supports_tools"] is True + assert flags["supports_reasoning"] is True + assert flags["supports_preserve_thinking"] is True diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py new file mode 100644 index 0000000000..923af87c4f --- /dev/null +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -0,0 +1,788 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Tests for the safetensors agentic tool loop. + +Covers the shared ``tool_call_parser`` helpers and the cumulative-text +state machine inside ``safetensors_agentic.run_safetensors_tool_loop``. +The loop is exercised with hand-crafted fake single-turn generators so +no model load is needed; the tests run in CI under a few seconds. + +Edge cases under coverage: +* Plain answers (no tool calls) flush full content. +* Single ``{json}`` triggers the tool and re-enters. +* Single ``...`` XML form triggers the same path. +* Truncated unclosed ```` is still parsed. +* Tool result is fed back as ``role=tool`` for the next iteration. +* Bad JSON inside ```` does not raise and (when healed) is + routed as a ``{"query": ...}`` web search call. +* Duplicate tool calls produce a synthetic "do not repeat" result the + second time. +* ``__IMAGES__`` sentinel is stripped before the model sees the result. +* Tool execution errors are tagged so the model gets a nudge but the + loop keeps streaming. +* Cancel is honoured between iterations. +* ``max_tool_iterations`` cap is respected and a final-answer attempt + closes the stream cleanly. +""" + +import threading + +import pytest + +from core.inference import safetensors_agentic +from core.inference.safetensors_agentic import ( + _coerce_arguments, + run_safetensors_tool_loop, +) +from core.inference.tool_call_parser import ( + has_tool_signal, + parse_tool_calls_from_text, + strip_tool_markup, +) +from utils.datasets import is_gpt_oss_model_name + + +# ──────────────────────────────────────────────────────────────────── +# parse_tool_calls_from_text +# ──────────────────────────────────────────────────────────────────── + + +class TestParser: + def test_json_tool_call(self): + text = ( + '{"name":"web_search","arguments":{"query":"hello"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + tc = result[0] + assert tc["type"] == "function" + assert tc["function"]["name"] == "web_search" + # Arguments must always be a JSON string. + assert isinstance(tc["function"]["arguments"], str) + assert "hello" in tc["function"]["arguments"] + + def test_json_tool_call_unclosed(self): + # No ; balanced-brace extractor must still close. + text = '{"name":"python","arguments":{"code":"print(1)"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_xml_function_call(self): + text = "print('hi')" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print('hi')" in result[0]["function"]["arguments"] + + def test_xml_unclosed(self): + # Closing tags omitted; parser must still extract the value. + text = "ls -la" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "terminal" + assert "ls -la" in result[0]["function"]["arguments"] + + def test_code_with_embedded_xml(self): + # A code parameter contains the literal . Must not + # truncate the value because the parser uses end-of-body as the + # only boundary for single-parameter calls. + text = ( + "html = ''\n" + "print('hi')" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert "print('hi')" in result[0]["function"]["arguments"] + + def test_multiple_calls(self): + text = ( + '{"name":"web_search","arguments":{"query":"a"}}' + '{"name":"web_search","arguments":{"query":"b"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "web_search" + assert result[1]["function"]["name"] == "web_search" + + def test_bad_json_does_not_raise(self): + text = "{not valid json}" + result = parse_tool_calls_from_text(text) + # Bad JSON is silently dropped; caller can fall back to text. + assert result == [] + + def test_has_tool_signal(self): + assert has_tool_signal("blah x") + assert has_tool_signal("hi ...") + assert not has_tool_signal("hello world") + + def test_strip_markup_closed(self): + text = "before {} after" + assert strip_tool_markup(text) == "before after" + + def test_strip_markup_unclosed_final(self): + text = "before {partial" + # With final=True the trailing run is dropped. + assert strip_tool_markup(text, final = True) == "before" + # Without final=True the unclosed run is preserved. + assert "partial" in strip_tool_markup(text) + + +# ──────────────────────────────────────────────────────────────────── +# run_safetensors_tool_loop +# ──────────────────────────────────────────────────────────────────── + + +def _fake_stream(chunks): + """Build a single-turn generator that yields cumulative snapshots.""" + + def _gen(_messages): + acc = "" + for c in chunks: + acc += c + yield acc + + return _gen + + +def _const_stream(text): + """A single-turn generator that yields one cumulative snapshot.""" + + def _gen(_messages): + yield text + + return _gen + + +class FakeExecuteTool: + """Stand-in for ``core.inference.tools.execute_tool``.""" + + def __init__(self, results): + # ``results`` is a list of strings or RuntimeError instances. + self.results = list(results) + self.calls: list[tuple[str, dict]] = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + ): + self.calls.append((name, arguments)) + result = self.results.pop(0) if self.results else "OK" + if isinstance(result, Exception): + raise result + return result + + +def _collect_events(generator, max_events = 200): + events = [] + for ev in generator: + events.append(ev) + if len(events) >= max_events: + break + return events + + +def _make_loop(*, turns, exec_results = None, **kwargs): + """Build a configured loop with a multi-turn fake generator. + + ``turns`` is a list of chunk-lists; iteration N yields chunks from + ``turns[N]``. + """ + turn_iter = iter(turns) + + def _gen(_messages): + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(exec_results or []) + return run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "terminal"}}, + ], + execute_tool = exec_fn, + **kwargs, + ), exec_fn + + +class TestLoopBasic: + def test_plain_answer(self): + # No tool XML; loop should yield content then status="". + loop, _exec = _make_loop( + turns = [["Hello", " world", "!"]], + exec_results = [], + ) + events = _collect_events(loop) + contents = [e for e in events if e["type"] == "content"] + statuses = [e for e in events if e["type"] == "status"] + assert contents, "expected at least one content event" + # Final cumulative content should contain the answer. + final_text = contents[-1]["text"] + assert "Hello world!" in final_text + assert statuses and statuses[-1]["text"] == "" + + def test_single_tool_then_answer(self): + loop, exec_fn = _make_loop( + turns = [ + # : tool call only. + [ + '{"name":"web_search",', + '"arguments":{"query":"weather"}}', + "", + ], + # : final answer. + ["The ", "weather is ", "sunny."], + ], + exec_results = ["Sunny and 22C"], + ) + events = _collect_events(loop) + kinds = [e["type"] for e in events] + + assert "tool_start" in kinds + assert "tool_end" in kinds + # Tool was actually called with the parsed arguments. + assert exec_fn.calls == [("web_search", {"query": "weather"})] + + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_name"] == "web_search" + tool_end = next(e for e in events if e["type"] == "tool_end") + assert tool_end["result"] == "Sunny and 22C" + + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + def test_function_xml_form(self): + loop, exec_fn = _make_loop( + turns = [ + ["print(1)"], + ["Result: 1"], + ], + exec_results = ["1\n"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("python", {"code": "print(1)"})] + contents = [e for e in events if e["type"] == "content"] + assert "Result: 1" in contents[-1]["text"] + + def test_truncated_unclosed_tool_call(self): + loop, exec_fn = _make_loop( + turns = [ + # No ; balanced-brace parser must still + # succeed because the JSON itself is balanced. + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["done"], + ], + exec_results = ["result"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + + def test_bad_json_healed_to_query(self): + # Tool call with non-JSON string arguments. With auto_heal_tool_calls + # the string is routed as {"query": ...}. + loop, exec_fn = _make_loop( + turns = [ + # JSON inside the tool call is well-formed; the + # ``arguments`` is a string that is not itself valid + # JSON for ``_coerce_arguments`` to parse, so the + # heal path runs. + [ + '{"name":"web_search","arguments":"hello world"}' + ], + ["ok"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls and exec_fn.calls[0][0] == "web_search" + assert exec_fn.calls[0][1] == {"query": "hello world"} + + +class TestLoopBehaviour: + def test_duplicate_tool_call_synthetic_result(self): + # Two identical successful calls in a row: the second is short- + # circuited with a "do not repeat" message and execute_tool is + # called only once. + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["final"], + ], + exec_results = ["search-result-1"], + ) + events = _collect_events(loop) + # Only one real call. + assert len(exec_fn.calls) == 1 + tool_end_events = [e for e in events if e["type"] == "tool_end"] + assert len(tool_end_events) == 2 + assert "do not repeat" in tool_end_events[1]["result"].lower() + + def test_image_sentinel_stripped_from_model_feed(self): + # The tool result has a frontend image sentinel that should be + # stripped before being fed back into the next turn, BUT the + # tool_end event still carries the raw result for the UI. + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"python","arguments":{"code":"plot()"}}' + ], + ["see chart"], + ], + exec_results = ["chart\n__IMAGES__:/tmp/chart.png"], + ) + events = _collect_events(loop) + tool_end = next(e for e in events if e["type"] == "tool_end") + assert "__IMAGES__" in tool_end["result"] + + def test_image_sentinel_stripped_with_leading_marker(self): + # Sentinel at start (no newline) must not leak to the model. + from core.inference import safetensors_agentic as _sa + + captured: list[list[dict]] = [] + + def fake_single_turn(messages, **_kw): + captured.append([dict(m) for m in messages]) + if len(captured) == 1: + yield '{"name":"python","arguments":{"code":"plot()"}}' + else: + yield "done" + + events = list( + _sa.run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "plot please"}], + tools = [{"function": {"name": "python"}}], + execute_tool = lambda *_a, **_kw: "__IMAGES__:/tmp/x.png", + cancel_event = threading.Event(), + max_tool_iterations = 3, + auto_heal_tool_calls = True, + ) + ) + # Model's second turn must not see "__IMAGES__". + assert len(captured) >= 2 + tool_msgs = [m for m in captured[1] if m.get("role") == "tool"] + assert tool_msgs, "no tool message reached the model" + for tm in tool_msgs: + assert ( + "__IMAGES__" not in tm["content"] + ), f"sentinel leaked to model: {tm['content']!r}" + + def test_image_sentinel_stripped_with_multiple_markers(self): + # Consecutive sentinels: cut at the first, nothing leaks. + from core.inference import safetensors_agentic as _sa + + captured: list[list[dict]] = [] + + def fake_single_turn(messages, **_kw): + captured.append([dict(m) for m in messages]) + if len(captured) == 1: + yield '{"name":"python","arguments":{"code":"plot()"}}' + else: + yield "done" + + multi = "panel\n__IMAGES__:/tmp/a.png\n__IMAGES__:/tmp/b.png" + events = list( + _sa.run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "plot please"}], + tools = [{"function": {"name": "python"}}], + execute_tool = lambda *_a, **_kw: multi, + cancel_event = threading.Event(), + max_tool_iterations = 3, + auto_heal_tool_calls = True, + ) + ) + tool_msgs = [m for m in captured[1] if m.get("role") == "tool"] + assert tool_msgs + for tm in tool_msgs: + assert ( + "__IMAGES__" not in tm["content"] + ), f"second sentinel leaked: {tm['content']!r}" + assert ( + tm["content"] == "panel" + ), f"expected payload-only 'panel', got {tm['content']!r}" + + def test_tool_execution_error_is_emitted_but_loop_continues(self): + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["sorry, that failed"], + ], + exec_results = ["Error: network unreachable"], + ) + events = _collect_events(loop) + tool_end = next(e for e in events if e["type"] == "tool_end") + assert tool_end["result"].startswith("Error") + # The loop must still produce a content event after the failure. + contents = [e for e in events if e["type"] == "content"] + assert contents + + def test_exception_in_executor_does_not_raise(self): + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["recovered"], + ], + exec_results = [RuntimeError("boom")], + ) + events = _collect_events(loop) + tool_end = next(e for e in events if e["type"] == "tool_end") + assert "boom" in tool_end["result"] + + +class TestLoopControl: + def test_cancel_event_breaks_loop(self): + cancel = threading.Event() + cancel.set() + # Even with a fake stream that emits tool calls, the loop must + # bail before invoking execute_tool when cancel is set. + exec_fn = FakeExecuteTool([]) + events = list( + run_safetensors_tool_loop( + single_turn = _const_stream( + '{"name":"web_search",' + '"arguments":{"query":"x"}}' + ), + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + cancel_event = cancel, + ) + ) + assert events == [] + assert exec_fn.calls == [] + + def test_max_iterations_caps_loop(self): + # The loop should stop after max_tool_iterations even if the + # model keeps asking for tools, then emit a final-attempt round. + loop, exec_fn = _make_loop( + turns = [ + # : tool call (executes once) + [ + '{"name":"web_search","arguments":{"query":"a"}}' + ], + # : model gives a final answer when nudged. + ["here is the final answer"], + ], + exec_results = ["result"], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e for e in events if e["type"] == "content"] + # Final content must include the final answer. + assert contents and "final answer" in contents[-1]["text"] + + +class TestStatusFormatting: + def test_status_for_known_tools(self): + # Use the private helper directly to verify status formatting. + assert ( + safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) + == "Searching: abc" + ) + assert ( + safetensors_agentic._status_for_tool( + "web_search", {"url": "https://www.example.com/x"} + ) + == "Reading: example.com" + ) + assert safetensors_agentic._status_for_tool( + "python", {"code": "x = 1"} + ).startswith("Running Python:") + assert safetensors_agentic._status_for_tool( + "terminal", {"command": "ls"} + ).startswith("Running:") + assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith( + "Calling:" + ) + + +class TestProseMentioningToolCall: + def test_assistant_prose_with_literal_tool_call_text_survives(self): + # Regression: if the assistant text legitimately mentions + # ```` as a literal string and the parser finds no + # actual call, the loop must surface the full content instead + # of silently stripping everything past the literal marker. + loop, exec_fn = _make_loop( + turns = [ + # : a real tool call so the loop moves to + # . + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + # : prose that mentions the literal text. + ["the docs say means an LLM tool call wrapper"], + ], + exec_results = ["result"], + ) + events = _collect_events(loop) + contents = [e for e in events if e["type"] == "content"] + assert contents, "expected at least one content event" + final = contents[-1]["text"] + assert ( + "LLM tool" in final + ), f"prose mentioning should not be truncated; got {final!r}" + + def test_tool_result_with_tool_call_text_does_not_retrigger(self): + # Tool result text contains the literal ```` string. + # The loop must only parse the MODEL output, not the tool + # result, so we should see exactly one call. + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["the docs mention wrappers"], + ], + exec_results = ["Page text: appears here in the docs"], + ) + events = _collect_events(loop) + assert len(exec_fn.calls) == 1 + + +class TestChatTemplateHelper: + """Cover the dependency-light helper used by InferenceBackend.""" + + def setup_method(self): + from core.inference.chat_template_helpers import ( + apply_chat_template_for_generation, + ) + + self.apply = apply_chat_template_for_generation + + class _Tok: + def __init__(self, accepted): + self.accepted = accepted + self.call_count = 0 + self.last_kwargs = None + + def apply_chat_template( + self, messages, *, tokenize = False, add_generation_prompt = True, **kw + ): + self.call_count += 1 + unknown = set(kw) - self.accepted + if unknown: + raise TypeError(f"unexpected kwargs: {sorted(unknown)}") + self.last_kwargs = dict(kw) + return "PROMPT" + + def test_richest_call_wins_when_template_supports_all(self): + tok = self._Tok({"tools", "enable_thinking"}) + self.apply(tok, [], tools = [{}], enable_thinking = True) + assert tok.call_count == 1 + assert "tools" in tok.last_kwargs + assert "enable_thinking" in tok.last_kwargs + + def test_falls_back_when_template_rejects_reasoning_kwarg(self): + tok = self._Tok({"tools"}) + self.apply(tok, [], tools = [{}], enable_thinking = True) + assert tok.call_count >= 2 + assert tok.last_kwargs == {"tools": [{}]} + + def test_falls_back_to_bare_call(self): + tok = self._Tok(set()) + self.apply(tok, [], tools = [{}], enable_thinking = True) + assert tok.last_kwargs == {} + + def test_jinja_error_propagates(self): + class Boom: + def apply_chat_template(self, *a, **kw): + raise ValueError("jinja: missing var") + + with pytest.raises(ValueError): + self.apply(Boom(), []) + + def test_no_kwargs_single_call(self): + tok = self._Tok(set()) + self.apply(tok, []) + assert tok.call_count == 1 + + +# ──────────────────────────────────────────────────────────────────── +# Guardrails (allowlist, budget, streaming-leak, dedup, id offset, +# auto_heal=False, canonical healed-arg key) +# ──────────────────────────────────────────────────────────────────── + + +class TestGuardrails: + def test_disabled_tool_is_not_executed(self): + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _fake_stream( + [ + '{"name":"terminal","arguments":{"command":"echo bypass"}}' + ] + ), + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert tool_ends and "not enabled" in tool_ends[0]["result"].lower() + + def test_empty_tools_list_does_not_enforce_allowlist(self): + exec_fn = FakeExecuteTool(["OK"]) + loop = run_safetensors_tool_loop( + single_turn = _fake_stream( + [ + '{"name":"python","arguments":{"code":"print(1)"}}' + ] + ), + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + _collect_events(loop) + assert exec_fn.calls == [("python", {"code": "print(1)"})] + + def test_max_iterations_zero_executes_no_tools(self): + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ] + ], + exec_results = ["OK"], + max_tool_iterations = 0, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert events and events[-1] == {"type": "status", "text": ""} + + def test_streaming_clips_before_tool_signal_no_leak(self): + loop, exec_fn = _make_loop( + turns = [ + [ + "I will look this up. ", + "Some more prose that's long enough to leave the buffer. ", + '{"name":"web_search","arguments":{"query":"x"}}', + ], + ["all done"], + ], + exec_results = ["weather: sunny"], + max_tool_iterations = 2, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + for e in events: + if e["type"] == "content": + assert "" not in e["text"] + assert "web_search" not in e["text"] + + def test_auto_heal_disabled_still_parses_valid_tool_call(self): + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["done"], + ], + exec_results = ["OK"], + auto_heal_tool_calls = False, + max_tool_iterations = 2, + ) + _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + + def test_non_consecutive_duplicate_is_short_circuited(self): + loop, exec_fn = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"A"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"B"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"A"}}' + ], + ["final"], + ], + exec_results = ["res-A", "res-B"], + max_tool_iterations = 4, + ) + events = _collect_events(loop) + assert exec_fn.calls == [ + ("web_search", {"query": "A"}), + ("web_search", {"query": "B"}), + ] + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert "already made this exact call" in tool_ends[-1]["result"] + + def test_coerce_string_args_python_uses_code_key(self): + assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == { + "code": "print(1)" + } + + def test_coerce_string_args_terminal_uses_command_key(self): + assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == { + "command": "ls -la" + } + + def test_tool_call_ids_unique_across_loop_iterations(self): + loop, _exec = _make_loop( + turns = [ + [ + '{"name":"web_search","arguments":{"query":"A"}}' + ], + [ + '{"name":"web_search","arguments":{"query":"B"}}' + ], + ["done"], + ], + exec_results = ["A", "B"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + ids = [e["tool_call_id"] for e in events if e["type"] == "tool_start"] + assert len(ids) == 2 and ids[0] != ids[1] + + +# ──────────────────────────────────────────────────────────────────── +# Shared gpt-oss name detector +# ──────────────────────────────────────────────────────────────────── + + +class TestGptOssNameDetection: + def test_substring_match(self): + assert is_gpt_oss_model_name("unsloth/gpt-oss-20b") is True + + def test_negative_known_non_oss_model(self): + assert is_gpt_oss_model_name("meta-llama/Llama-3.1-8B-Instruct") is False + + def test_empty_or_none_returns_false(self): + assert is_gpt_oss_model_name("") is False + assert is_gpt_oss_model_name(None) is False + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py index c9237d83c1..7988b09972 100644 --- a/studio/backend/utils/datasets/__init__.py +++ b/studio/backend/utils/datasets/__init__.py @@ -59,6 +59,7 @@ from .model_mappings import ( TEMPLATE_TO_MODEL_MAPPER, MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER, + is_gpt_oss_model_name, ) # Legacy imports from the original dataset_utils.py for backward compatibility @@ -98,6 +99,7 @@ __all__ = [ "TEMPLATE_TO_MODEL_MAPPER", "MODEL_TO_TEMPLATE_MAPPER", "TEMPLATE_TO_RESPONSES_MAPPER", + "is_gpt_oss_model_name", # Main entry points "check_dataset_format", "format_and_template_dataset", diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 21e8566ac5..eb2e5482c9 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -442,6 +442,26 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items(): MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key +def is_gpt_oss_model_name(name: str) -> bool: + """Name-based check for gpt-oss / harmony models. + + Used by both the in-process backend and the parent-process + orchestrator to detect harmony models without an IPC round-trip. + """ + name = (name or "").lower() + if not name: + return False + try: + if MODEL_TO_TEMPLATE_MAPPER.get(name) == "gpt-oss": + return True + for key, tmpl in MODEL_TO_TEMPLATE_MAPPER.items(): + if tmpl == "gpt-oss" and (key in name or name in key): + return True + except Exception: + pass + return "gpt-oss" in name + + TEMPLATE_TO_RESPONSES_MAPPER = { "gemma-4-thinking": { "instruction": "<|turn>user\n", diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index f87229d78f..359099e8b3 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -601,6 +601,12 @@ async function autoLoadSmallestModel(): Promise<{ reasoningStyle: sfLoadResp.reasoning_style ?? "enable_thinking", supportsPreserveThinking: sfLoadResp.supports_preserve_thinking ?? false, supportsTools: sfLoadResp.supports_tools ?? false, + // Parity with the GGUF branch above. + toolsEnabled: sfLoadResp.supports_tools ?? false, + codeToolsEnabled: sfLoadResp.supports_tools ?? false, + defaultChatTemplate: sfLoadResp.chat_template ?? null, + chatTemplateOverride: null, + loadedChatTemplateOverride: null, }); const sfModel: ChatModelSummary = { id: repo.repo_id, From 8d2440544006a46a8aa3fe49445fc02f7e599c0b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:40:25 -0700 Subject: [PATCH 20/32] studio/frontend: widen settings sidebar so 'Connections' label fits (#5607) The settings dialog sidebar was fixed at w-[200px], which left only ~92px of horizontal space for tab labels after icon, gap, and the 'New' badge for Connections/API. 'Connections' (11 chars at the 14.5px font weight medium) overflowed and rendered as 'Connectio...', matching the paper-cut reported in issue #5572. Bump the sidebar to w-[216px] -- 16 more pixels of label space, fully within the existing dialog width and unchanged on mobile (max-sm:w-full still drives the responsive layout). Co-authored-by: danielhanchen --- studio/frontend/src/features/settings/settings-dialog.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 1ece504d6f..e52257735e 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -110,7 +110,7 @@ export function SettingsDialog() { Manage your Unsloth Studio preferences.
-
)} - {st?.predicted_ms != null && ( + {hasPredicted && st?.predicted_ms != null && (
Generation @@ -93,11 +106,11 @@ export const MessageTiming: FC<{
)} - {st?.predicted_per_second != null && ( + {predictedRate != null && (
Speed - {st.predicted_per_second.toFixed(1)} tok/s + {predictedRate.toFixed(1)} tok/s
)} From 3b08d4a431c70001aa170b6bcca6fe21a226ee41 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:10 -0700 Subject: [PATCH 23/32] studio/web: differentiate offline from backend-down in fetch error (#5591) * studio/web: distinguish "offline" from "studio crashed" in error toast When the user's browser loses network mid-request, authFetch caught the fetch TypeError and surfaced "Studio isn't running -- please relaunch it." That is a correct diagnosis in the Tauri desktop app (the supervisor died in-process), but it is a misleading diagnosis in the web build where the backend lives elsewhere: the user will start hunting for a dead process when the actual problem is connectivity. Branch on navigator.onLine === false (web build only) and surface "You appear to be offline. Check your network connection and try again." instead. Tauri keeps the original wording so it stays accurate there. Found while running a slow-network UX probe and toggling Network.emulateNetworkConditions {offline: true} mid-stream. * studio/frontend: shorten offline-error wording comment --------- Co-authored-by: danielhanchen --- studio/frontend/src/features/auth/api.ts | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/studio/frontend/src/features/auth/api.ts b/studio/frontend/src/features/auth/api.ts index cc5749a923..1c0909454d 100644 --- a/studio/frontend/src/features/auth/api.ts +++ b/studio/frontend/src/features/auth/api.ts @@ -159,6 +159,14 @@ export async function authFetch( }); } catch (err) { if (err instanceof TypeError) { + // fetch TypeError = offline | backend down | CORS/DNS. In Tauri + // it's always backend-down; in the web build distinguish offline + // so the user gets the right recovery path. + if (!isTauri && typeof navigator !== "undefined" && navigator.onLine === false) { + throw new Error( + "You appear to be offline. Check your network connection and try again.", + ); + } throw new Error("Studio isn't running -- please relaunch it."); } throw err; From 0fb86e15b9eb3ddb073cec5a187adfef066af1d2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:15 -0700 Subject: [PATCH 24/32] studio/frontend: keep theme classes mutually exclusive on (#5580) * studio/frontend: keep theme classes mutually exclusive on The Sonner Toaster reads next-themes (mounted at provider.tsx with attribute="class" defaultTheme="light"), so on first mount next-themes adds a "light" class to . Studio's own setTheme path (features/settings/stores/theme-store.ts) only toggled "dark", so after the user picked Dark in settings the document ended up with html.className = "light dark". Harmless in CSS cascade because the dark variables override, but reads as a UI defect in devtools and trips CSS-aware tooling that branches on class lists. Toggle "light" alongside "dark" in applyToDocument so the two classes stay mutually exclusive regardless of how next-themes seeded the initial class. * studio/frontend: shorten theme-toggle comment --------- Co-authored-by: danielhanchen --- .../frontend/src/features/settings/stores/theme-store.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/features/settings/stores/theme-store.ts b/studio/frontend/src/features/settings/stores/theme-store.ts index 1879459a76..f73fc51e1c 100644 --- a/studio/frontend/src/features/settings/stores/theme-store.ts +++ b/studio/frontend/src/features/settings/stores/theme-store.ts @@ -32,7 +32,12 @@ function resolveTheme(theme: Theme): ResolvedTheme { function applyToDocument(resolved: ResolvedTheme) { if (typeof document === "undefined") return; - document.documentElement.classList.toggle("dark", resolved === "dark"); + // Keep "dark"/"light" mutually exclusive. next-themes (via Sonner) + // adds "light" on first mount; without the explicit toggle we'd end + // up with `class="light dark"` after a switch. + const cl = document.documentElement.classList; + cl.toggle("dark", resolved === "dark"); + cl.toggle("light", resolved === "light"); } const listeners = new Set<() => void>(); From 7b9fcf8cbc8e0ae160938d5160ba722dd0308cf9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:20 -0700 Subject: [PATCH 25/32] studio/frontend: show Loading fallback instead of blank pane on lazy route navigation (#5568) * studio/frontend: show Loading fallback instead of blank pane on lazy route navigation Closes #5567. Train, Recipes and Export pages are imported via React.lazy() in their respective createRoute calls, and the Suspense boundary around in __root.tsx passes fallback={null}. The result is a 1-3 second completely white pane between sidebar click and content paint, which is the exact failure mode behind reports that those pages look broken or stuck. /chat does not suffer from this because chat.tsx imports its ChatPage synchronously. Replace fallback={null} on both Suspense boundaries (hideNavbar and sidebar layouts) with a small centered 'Loading...' label using the same muted-foreground style as elsewhere in the app. Synchronous routes (/chat) never suspend so they are unaffected; lazy routes now have a visible terminal-state placeholder while their chunk loads. * studio/frontend: also apply RouteFallback to the sidebar Suspense The first revision only replaced the fallback={null} inside the hideNavbar branch (used for onboarding / login). The primary lazy boundary that wraps Train / Recipes / Export is inside the SidebarInset branch at the other Suspense site, which kept rendering null and made the page look stuck for the same window the original bug describes (per bot review feedback on #5568). Replace both Suspense fallbacks with RouteFallback so the "Loading..." placeholder fires on every lazy route, not just on the auth flows. * studio/frontend: shorten RouteFallback comment --------- Co-authored-by: danielhanchen --- studio/frontend/src/app/routes/__root.tsx | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 22d149473c..295f041129 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -15,9 +15,17 @@ import { useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; -import { Suspense, useEffect } from "react"; +import { Suspense, useEffect, type ReactNode } from "react"; import { AppProvider } from "../provider"; +// Fallback while a lazy route bundle (Train/Recipes/Export) loads. +// /chat is synchronous and never hits this. +const RouteFallback: ReactNode = ( +
+ Loading... +
+); + const CHAT_ONLY_ALLOWED = new Set([ "/", "/chat", @@ -72,7 +80,7 @@ function RootLayout() { {hideNavbar ? (
- +
@@ -98,7 +106,7 @@ function RootLayout() { transition={{ duration: 0.15 }} className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"}`} > - + From 75ee380a07d92e3c7be73ccd2aacfbf1d64bad10 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:25 -0700 Subject: [PATCH 26/32] studio/frontend: include filename in attachment aria-label + img alt (#5594) * studio/frontend: include filename in attachment aria-label and img alt When a chat has multiple attachments of the same kind, the rendered tiles all share the generic accessible name "Image attachment" or "Document attachment". Sighted users get the filename from the Radix tooltip that pops on hover, but: - screen-reader users hear "Image attachment, Image attachment, Image attachment" with no way to distinguish three PNGs; - touch-device users (no hover) lose the filename entirely; - keyboard-only users would have to focus and read a tooltip that isn't always announced. Fold the filename into both the button's aria-label and the thumbnail , falling back to the existing labels when the attachment has no filename. Sighted UX is unchanged: the Radix tooltip already shows the same name on hover, and the visible aria-label has no rendered counterpart. Found while running a multi-image attach probe in the autonomous Studio UX loop (cycle 8). Repro: await page.evaluate(`Array.from(document.querySelectorAll( 'button[aria-label*="attachment" i]' )).map(b => b.getAttribute('aria-label'))`) Before: ["Image attachment", "Document attachment", "Add Attachment"] After: ["Image attachment: test_red_circle.png", "Document attachment: notes.txt", "Add Attachment"] * studio/frontend: shorten attachment a11y comment --------- Co-authored-by: danielhanchen --- .../src/components/assistant-ui/attachment.tsx | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx index b5b2810008..3ae1c68561 100644 --- a/studio/frontend/src/components/assistant-ui/attachment.tsx +++ b/studio/frontend/src/components/assistant-ui/attachment.tsx @@ -120,12 +120,13 @@ const AttachmentPreviewDialog: FC = ({ children }) => { const AttachmentThumb: FC = () => { const src = useAttachmentSrc(); + const name = useAuiState(({ attachment }) => attachment.name); if (src) { return ( Attachment preview ); @@ -143,6 +144,7 @@ const AttachmentUI: FC = () => { const isComposer = aui.attachment.source === "composer"; const isImage = useAuiState(({ attachment }) => attachment.type === "image"); + const name = useAuiState(({ attachment }) => attachment.name); const typeLabel = useAuiState(({ attachment }) => { const type = attachment.type; switch (type) { @@ -156,6 +158,11 @@ const AttachmentUI: FC = () => { throw new Error(`Unknown attachment type: ${type as string}`); } }); + // Include filename in accessible name so screen readers distinguish + // same-typed attachments. Sighted users get it via the tooltip. + const accessibleName = name + ? `${typeLabel} attachment: ${name}` + : `${typeLabel} attachment`; return ( @@ -175,7 +182,7 @@ const AttachmentUI: FC = () => { "aui-attachment-tile-composer border-foreground/20", )} id="attachment-tile" - aria-label={`${typeLabel} attachment`} + aria-label={accessibleName} type="button" > From feadfd5c1b1e8ac0dae42bad5a9c0ab5fa68a1ce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:30 -0700 Subject: [PATCH 27/32] studio/frontend: compare composer blocks send when no model picked (#5574) * studio/frontend: compare composer blocks send when no model picked Closes the racing-handle half of #5569. In Compare mode (GeneralCompare shell with model1/model2 props), if the user sends a prompt before picking models in either pane, the SharedComposer used to fall through to the per-handle append branch. Both panes then raced createOpenAIStreamAdapter -> autoLoadSmallestModel, one won, the other dispatched into an unloaded slot and produced an empty bubble with a 1000000.0 tok/s readout. The per-pane picker state never observed the global checkpoint change either, so both pickers stayed at "Select model". Add a guard before the content build: when handlesRef has model1/model2 keys but both selections are empty, surface a toast asking the user to pick models first, leave the text in the composer for retry, and never enter the racing dispatch path. Keeps the per-pane picker state as the source of truth for which model is on each side. The unphysical tok/s readout that the same path produced is separately covered by PR #5570 (display guard). * studio/frontend: tighten compare-mode guard to require both panes Review feedback on #5574: - Gemini: the redundant `model1 !== undefined && model2 !== undefined` checks let the racing-handle dispatch slip through whenever the Compare props arrive as undefined, which is the exact case the guard is trying to block. - Codex: with `isGeneralizedCompare` keyed on `model1?.id || model2?.id`, a half-selected Compare (one model picked, one empty) still falls into the generalized branch. The composer clears, the empty pane gets the user message appended, and `startRun` only fires for the side with an id, leaving the empty pane with a dangling prompt and no response. Switch `isGeneralizedCompare` to require BOTH panes (`&&`), drop the undefined gate, and surface the "Pick a model in each pane" toast for either the fully-empty or half-selected case. `hasCompareHandles` is true only inside GeneralCompareContent, so LoraCompare and the single-pane path stay unchanged. * studio/frontend: shorten compare-mode no-model-guard comment * studio/frontend: clarify compare-pane toast wording --------- Co-authored-by: danielhanchen --- .../frontend/src/features/chat/shared-composer.tsx | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index aef004e891..11a73ee486 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -523,7 +523,19 @@ export function SharedComposer({ handlesRef.current["model1"] || handlesRef.current["model2"], ); const isGeneralizedCompare = - hasCompareHandles && Boolean(model1?.id || model2?.id); + hasCompareHandles && Boolean(model1?.id && model2?.id); + + // Generalized compare requires both panes to have a model. A + // half-selected send either races to an empty bubble with bogus + // tok/s (#5569) or leaves the empty pane with a dangling prompt. + // hasCompareHandles is true only in GeneralCompareContent, so + // LoraCompare and single-pane chats are unaffected. + if (hasCompareHandles && !isGeneralizedCompare) { + toast.error("Pick a model in each pane to compare", { + description: "Use the model dropdown above each pane, then send your prompt.", + }); + return; + } if (pendingImages.length > 0 && !isGeneralizedCompare && imageUnavailableReason) { // Single mode: the loaded model's runtime capability is known From cf53ff68616bdf1359bfc39304157d2b228f3175 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:48 -0700 Subject: [PATCH 28/32] studio: restore focus to opener when settings dialog closes (#5612) The settings dialog opens via a global Ctrl+, keydown handler in __root.tsx, not via a . Radix's FocusScope tries to capture document.activeElement at mount as the focus-restore target, but settings-dialog.tsx schedules a requestAnimationFrame that focuses the active tab button right after mount, racing FocusScope's previous- focus capture. On Escape or close-button click, focus then lands on instead of the textarea (or button, or wherever the user was). A Playwright focus-management probe confirmed: open dialog, press Tab 15 times (trap holds), press Escape, document.activeElement === BODY. This is a WCAG 2.4.3 (Focus Order) violation: keyboard-only users have to re-Tab from the start of the page after every settings visit. Fix: capture document.activeElement in the Zustand store at the moment openDialog() runs, then restore via onCloseAutoFocus on DialogContent. Use opener.isConnected so a stale node from a re-rendered tree falls back to Radix's default. closeDialog deliberately does NOT clear the opener slot - onCloseAutoFocus reads it on the render after open=false, so clearing in the same set() would null it before restoration. Probe re-run confirms focus restored to the TEXTAREA opener after Escape, after close-button click, on both repeats. Tab + Shift+Tab trap still holds (unchanged Radix behaviour). Co-authored-by: danielhanchen --- .../src/features/settings/settings-dialog.tsx | 10 ++++++++++ .../settings/stores/settings-dialog-store.ts | 15 +++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index e52257735e..384e359f31 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -74,6 +74,7 @@ export function SettingsDialog() { const activeTab = useSettingsDialogStore((s) => s.activeTab); const setActiveTab = useSettingsDialogStore((s) => s.setActiveTab); const closeDialog = useSettingsDialogStore((s) => s.closeDialog); + const opener = useSettingsDialogStore((s) => s.opener); const reduced = useReducedMotion(); const tabButtonRefs = useRef>({ general: null, @@ -98,6 +99,15 @@ export function SettingsDialog() { { + // Restore focus to the element that triggered openDialog(). + // Radix's FocusScope races our rAF-scheduled tab-button focus + // and loses the previous-focus reference, so we restore by hand. + if (opener && opener.isConnected) { + e.preventDefault(); + opener.focus({ preventScroll: true }); + } + }} className={cn( "!max-w-none h-[560px] w-[820px] p-0 overflow-hidden", "shadow-border rounded-xl border-border", diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts index a2e32f26a3..7ac5422ff1 100644 --- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts +++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts @@ -15,6 +15,11 @@ export type SettingsTab = interface SettingsDialogState { open: boolean; activeTab: SettingsTab; + // Element focused at the moment openDialog() ran. Radix's FocusScope + // would normally track this, but the rAF-scheduled focus() in + // settings-dialog.tsx races its previous-focus capture, leaving focus + // on after close. We restore explicitly via onCloseAutoFocus. + opener: HTMLElement | null; openDialog: (tab?: SettingsTab) => void; closeDialog: () => void; setActiveTab: (tab: SettingsTab) => void; @@ -47,11 +52,21 @@ function loadInitialTab(): SettingsTab { export const useSettingsDialogStore = create((set) => ({ open: false, activeTab: loadInitialTab(), + opener: null, openDialog: (tab) => set((state) => ({ open: true, activeTab: tab ?? state.activeTab, + opener: + typeof document !== "undefined" && + document.activeElement instanceof HTMLElement && + document.activeElement !== document.body + ? document.activeElement + : null, })), + // Do NOT clear `opener` here. onCloseAutoFocus runs on the next render + // pass after `open: false` lands, so the opener must still be readable + // from the store at that point. The next openDialog() overwrites it. closeDialog: () => set({ open: false }), setActiveTab: (tab) => { try { From 7d84499197fb6cbace1b7b23f5927292fead6ce6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:56:53 -0700 Subject: [PATCH 29/32] studio/frontend: add aria-label to Dictate / Stop dictation buttons (#5599) The composer's mic icon buttons used tooltip="Dictate" / "Stop dictation" but no aria-label, so screen-reader users heard only the empty SVG-only button. Every other composer icon button (Send, Add Attachment, audio buttons, composer pills) carries an explicit aria-label; the shared-composer.tsx implementation already does too. Mirror that here for parity. Co-authored-by: danielhanchen --- studio/frontend/src/components/assistant-ui/thread.tsx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index d72547aa1a..8776d5d6d6 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -916,6 +916,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({ @@ -927,6 +928,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({ From 3dbddc39c2b92f46560a4c9ce4d9d87b344d7a51 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:57:08 -0700 Subject: [PATCH 30/32] studio/frontend: settings dialog fits viewport at tablet widths (#5600) * studio/frontend: settings dialog fits viewport at tablet widths The dialog used a fixed w-[820px] with sm:w-[820px] override, so any viewport between 640px and 820px (iPad portrait at 768px is the canonical case) saw the dialog overflow horizontally by 26px on each side -- the right-edge scroll arrow and the active-tab chevron got clipped against the viewport. Replace the hard 820 with min(820px, calc(100vw-2rem)) on both max-w and w so the dialog caps at the original 820px on desktop and shrinks to fit (with a 1rem gutter) on narrower screens. max-sm: still drives the full-bleed h-dvh/w-dvw layout under 640px. * studio/frontend: keep mobile full-bleed override !important Bot review: base !max-w-[min(...)] is !important so the regular max-sm:max-w-none never wins, leaving a 1rem gutter on phones where the previous code rendered a true full-bleed dialog. Bump the mobile override to !important too. --------- Co-authored-by: danielhanchen --- studio/frontend/src/features/settings/settings-dialog.tsx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 384e359f31..3b95cb2355 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -109,10 +109,12 @@ export function SettingsDialog() { } }} className={cn( - "!max-w-none h-[560px] w-[820px] p-0 overflow-hidden", + // Cap at 820px but shrink to the viewport so we don't clip + // on iPad-portrait widths (640-820px) where the fixed + // `w-[820px]` overflows by 26px on each side. + "!max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden", "shadow-border rounded-xl border-border", - "sm:h-[560px] sm:w-[820px]", - "max-sm:h-dvh max-sm:w-dvw max-sm:rounded-none", + "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", )} > Settings From ebed6469a023be167c588fc82482a123d3eb0dfa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 06:57:14 -0700 Subject: [PATCH 31/32] studio/frontend: show Generation stopped placeholder when cancelled mid-thinking (#5565) * studio/frontend: show Generation stopped placeholder when cancelled mid-thinking Closes #5563. When the user clicks Stop before any visible content has streamed in, the running indicator disappears but no Parts have rendered yet, leaving just the AssistantActionBar floating below the user prompt. That looks broken (and is the exact failure mode behind the 'tools work, but I don't see anything happening' bucket of reports). Add a sibling CancelledIndicator next to GeneratingIndicator that fires when content is empty AND status is incomplete with reason cancelled, rendering a muted 'Generation stopped.' italic. The terminal-state label is consistent with tool-fallback's existing 'Cancelled tool' treatment and with reasoning's 'Thought for N seconds' summary. * studio/frontend: shorten CancelledIndicator comment Trim the 3-line explanation to a single line describing what the placeholder is for. * studio/frontend: use 'Cancelled.' to match tool-fallback wording --------- Co-authored-by: danielhanchen --- .../src/components/assistant-ui/thread.tsx | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 8776d5d6d6..0326b90a97 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -996,6 +996,24 @@ const GeneratingIndicator: FC = () => { return Generating...; }; +// Placeholder when stop fires before any visible content (e.g. mid-think). +const CancelledIndicator: FC = () => { + const show = useAuiState( + ({ message }) => + message.content.length === 0 && + message.status?.type === "incomplete" && + message.status?.reason === "cancelled", + ); + if (!show) { + return null; + } + return ( + + Cancelled. + + ); +}; + const AssistantMessage: FC = () => { return ( { >
+ Date: Tue, 19 May 2026 07:00:04 -0700 Subject: [PATCH 32/32] 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 c99a182ce6..07dd9dd8f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.5.2", + "unsloth_zoo>=2026.5.3", "wheel>=0.42.0", "packaging", "numpy", @@ -90,7 +90,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.5.2", + "unsloth_zoo>=2026.5.3", "torchvision", "unsloth[triton]", ] @@ -580,7 +580,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.5.2", + "unsloth_zoo>=2026.5.3", "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 f4313c4518..572ec8842a 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.5.4" +__version__ = "2026.5.5" __all__ = [ "SUPPORTS_BFLOAT16",