From dfb3eedf77d1e04ce68a5b2de0aa9f7a3c81acc0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 23 May 2026 21:48:12 -0700 Subject: [PATCH 01/43] ci: broaden Linux + narrow Windows llama.cpp runtime patterns + trim #5741 comments (#5746) * ci: broaden Linux llama.cpp runtime pattern to lib*.so* #5741 patched the explicit Linux pattern list to add ``libllama-*-impl.so*`` after ggml-org/llama.cpp#23462 (between b9279 and b9283) split each binary's entry code into a paired ``lib-impl.so`` shared library. Same class of upstream repackaging will hit us again whenever a new shared lib is added. Mirror what macOS already does and replace the per-lib list with a single ``lib*.so*`` glob. ``copy_globs`` (line 3614) unions patterns, so the per-variant ``libggml-cuda.so*`` / ``libggml-hip.so*`` entries were never filtering anything; the spec lives in ``runtime_payload_health_groups`` (line 5209) which keeps the explicit minimum-required list per variant. Dry-run against b9296-bin-ubuntu-x64.tar.gz: 40 files copied (all ggml, llama, mtmd, impl variants + the two binaries we ship), 22 skipped (other CLIs, rpc-server, LICENSE). Functionally equal to the post-#5741 set. * cleanup: trim #5741 comments on the pydantic split Comments added in #5741 explained the original bug in full each time. They are mostly redundant with the commit message and the PR. Trim them to one short paragraph per site. No behavior change. * ci: narrow Windows runtime pattern to llama-server.exe + llama-quantize.exe Studio only invokes llama-server and llama-quantize. Mac and Linux already filter to those two binaries; Windows was the odd one out with ``*.exe`` copying every CLI upstream ships (llama-cli, llama-bench, llama-mtmd-cli, ...). Dry-run on b9296 (win cpu-x64, cpu-arm64, cuda-13.1, hip-radeon): 20 unused EXEs skipped per variant, all DLLs (incl. the new llama-*-impl.dll family) still copied via ``*.dll``. ``existing_install_matches_choice`` already checks llama-server.exe exists explicitly (line 5297), so the health gate is unchanged. --- install.ps1 | 13 +++---- install.sh | 13 +++---- .../backend/requirements/no-torch-runtime.txt | 18 +++------- studio/install_llama_prebuilt.py | 35 +++++-------------- studio/install_python_stack.py | 21 +++-------- tests/studio/install/test_rocm_support.py | 23 ++++++------ 6 files changed, 38 insertions(+), 85 deletions(-) diff --git a/install.ps1 b/install.ps1 index b26566cc3d..07512de720 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1302,13 +1302,9 @@ shell.Run cmd, 0, False # 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.6" unsloth-zoo } if ($baseInstallExit -eq 0) { - # Install pydantic WITH deps so pip pins pydantic-core to - # the exact version pydantic's metadata requires. The - # --no-deps install of no-torch-runtime.txt below would - # otherwise pick the latest of each independently and - # trip pydantic's _ensure_pydantic_core_version check. - # pydantic's deps (annotated-types, pydantic-core, - # typing-extensions, typing-inspection) are torch-free. + # Resolve pydantic WITH deps so pip pins pydantic-core + # to the matching version (no-torch-runtime.txt below + # is --no-deps). All transitive deps are torch-free. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { @@ -1358,8 +1354,7 @@ shell.Run cmd, 0, False # 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.6" unsloth-zoo } if ($baseInstallExit -eq 0) { - # Install pydantic WITH deps so pip pins pydantic-core to - # the matching version (see migrated branch above). + # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { diff --git a/install.sh b/install.sh index 9bdd935171..1e19e951de 100755 --- a/install.sh +++ b/install.sh @@ -1866,13 +1866,9 @@ if [ "$_MIGRATED" = true ]; then 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.6" unsloth-zoo - # Install pydantic WITH deps so pip pins pydantic-core to the - # exact version pydantic's own metadata requires. The --no-deps - # install below would otherwise pick the latest of each - # independently and trip pydantic's _ensure_pydantic_core_version - # check on the next import. pydantic's deps (annotated-types, - # pydantic-core, typing-extensions, typing-inspection) are - # torch-free, so this is safe on the no-torch path. + # Resolve pydantic WITH deps so pip pins pydantic-core to the + # matching version (no-torch-runtime.txt below is --no-deps). + # All transitive deps are torch-free. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" @@ -2051,8 +2047,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then 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.6" unsloth-zoo - # Install pydantic WITH deps so pip pins pydantic-core to the - # exact version pydantic requires (see migrated branch above). + # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index c33ebf4d94..85294114b1 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -22,19 +22,11 @@ rich>=13.0 markdown-it-py>=3.0 mdurl>=0.1 pygments>=2.0 -# pydantic is intentionally NOT installed via this --no-deps file. -# install.sh / install.ps1 / install_python_stack.py run a separate -# `pip install pydantic` (with deps) just before this file is -# applied, so pip resolves `pydantic-core` to the exact version -# pydantic's `_ensure_pydantic_core_version` check expects. Listing -# pydantic + pydantic-core unpinned here and resolving them under -# --no-deps used to pick the latest of each independently and trip -# `SystemError: pydantic-core 2.X.Y is incompatible with the current -# pydantic version` on the first import (Windows fresh-venv repro -# was the canonical case). pydantic's transitive deps -# (annotated-types, pydantic-core, typing-extensions, -# typing-inspection) are torch-free, so installing it WITH deps -# does not pull torch. +# pydantic is intentionally NOT pinned here. install.sh / install.ps1 +# / install_python_stack.py run `pip install pydantic` WITH deps just +# before this --no-deps file is applied, so pip resolves pydantic-core +# to the exact version pydantic's _ensure_pydantic_core_version check +# expects. Pinning both under --no-deps used to drift them apart. pyyaml nest-asyncio diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 91076a4743..7672af6630 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -3827,37 +3827,18 @@ def paired_runtime_dll_patterns(choice: AssetChoice) -> list[str]: def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: + # Broad shared-library glob + explicit binary names. Lets upstream + # repackage the SO/DLL set (e.g. ggml-org/llama.cpp#23462 split the + # per-binary entry code into paired ``lib-impl.so`` shared + # libraries between b9279 and b9283) without us re-enumerating + # every new file. Studio only invokes llama-server and llama-quantize; + # other CLIs upstream ships (llama-cli, llama-bench, ...) are skipped. if choice.install_kind in {"linux-cpu", "linux-cuda", "linux-rocm"}: - return [ - "llama-server", - "llama-quantize", - "libllama-common.so*", - "libllama.so*", - # Upstream llama.cpp split the per-binary entry code into - # paired ``libllama--impl.so`` shared libraries - # around release b9261. ``llama-server`` and - # ``llama-quantize`` are NEEDED-linked against - # ``libllama-server-impl.so`` / ``libllama-quantize-impl.so`` - # respectively, with RUNPATH ``$ORIGIN``. Without copying - # the impl ``.so`` files alongside the binaries, ldd - # reports them missing, preflight rejects the install, and - # the installer falls back to a source build on a fresh - # Linux install. Glob the whole family so future bundles - # that split additional binaries (e.g. ``llama-cli``, - # ``llama-bench``) keep working. - "libllama-*-impl.so*", - "libggml.so*", - "libggml-base.so*", - "libmtmd.so*", - "libggml-cpu-*.so*", - "libggml-cuda.so*", - "libggml-hip.so*", - "libggml-rpc.so*", - ] + return ["llama-server", "llama-quantize", "lib*.so*"] if choice.install_kind in {"macos-arm64", "macos-x64"}: return ["llama-server", "llama-quantize", "lib*.dylib"] if choice.install_kind in {"windows-cpu", "windows-cuda", "windows-hip"}: - return ["*.exe", "*.dll"] + return ["llama-server.exe", "llama-quantize.exe", "*.dll"] raise PrebuiltFallback( f"unsupported install kind for runtime overlay: {choice.install_kind}" ) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 4dfa20032b..9166d35ce3 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -979,22 +979,11 @@ def install_python_stack() -> int: package_name, "unsloth-zoo", ) - # Pydantic ships its core as a separate compiled wheel - # (pydantic-core), and pydantic's ``_ensure_pydantic_core_version`` - # checks the installed core matches the exact version pinned in - # its own metadata. With ``--no-deps`` plus an unpinned - # ``pydantic`` / ``pydantic-core`` pair in no-torch-runtime.txt, - # pip resolved each to the newest available version and the two - # drifted (pydantic 2.13.4 pins pydantic-core==2.46.4 today, but - # pydantic-core 2.47.0 was the latest). On a fresh Windows venv - # the next ``import pydantic`` raised ``SystemError: ... - # incompatible with the current pydantic version``. - # - # Resolve them WITH deps in a focused pip call so pip picks a - # compatible pair. pydantic's own deps are - # ``annotated-types``, ``pydantic-core``, ``typing-extensions``, - # ``typing-inspection`` -- none of which transitively pull - # torch, so this is safe for the no-torch path. + # Resolve pydantic WITH deps so pip pins pydantic-core to the + # exact version pydantic's metadata declares. Under --no-deps + # alone pip picks the latest of each and trips pydantic's + # _ensure_pydantic_core_version check. Transitive deps are + # torch-free. pip_install( "Installing pydantic (with deps for compatible core)", "--no-cache-dir", diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 698625e59b..e6f1ae1c65 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -346,28 +346,26 @@ class TestRuntimePatterns: patterns = runtime_patterns_for_choice(choice) assert "llama-server" in patterns assert "llama-quantize" in patterns - # Upstream split entry code into ``libllama--impl.so`` - # shared libraries (b9261+). llama-server and llama-quantize - # are NEEDED-linked against ``libllama-server-impl.so`` and - # ``libllama-quantize-impl.so`` respectively with RUNPATH - # ``$ORIGIN``, so the prebuilt overlay MUST copy them - # alongside the binaries or ldd reports them missing and - # preflight forces a source-build fallback. - assert "libllama-*-impl.so*" in patterns + # Broad lib*.so* covers libllama, libggml, libmtmd, libggml-cpu-*, + # plus the libllama--impl.so split that ggml-org/llama.cpp + # #23462 introduced between b9279 and b9283. + assert "lib*.so*" in patterns def test_linux_cuda_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cuda" ) patterns = runtime_patterns_for_choice(choice) - assert "libggml-cuda.so*" in patterns + # libggml-cuda.so is matched by lib*.so* now. + assert "lib*.so*" in patterns def test_linux_rocm_patterns(self): choice = AssetChoice( repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-rocm" ) patterns = runtime_patterns_for_choice(choice) - assert "libggml-hip.so*" in patterns + # libggml-hip.so is matched by lib*.so* now. + assert "lib*.so*" in patterns assert "llama-server" in patterns def test_windows_hip_patterns(self): @@ -380,7 +378,10 @@ class TestRuntimePatterns: install_kind = "windows-hip", ) patterns = runtime_patterns_for_choice(choice) - assert "*.exe" in patterns + # Narrowed from "*.exe" to the two binaries Studio actually + # invokes, mirroring the Linux/macOS pattern style. + assert "llama-server.exe" in patterns + assert "llama-quantize.exe" in patterns assert "*.dll" in patterns def test_macos_patterns(self): From 56e9046b2ffd9d6e54d9cd25d4f37d7f24f7ed4e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 23 May 2026 22:28:59 -0700 Subject: [PATCH 02/43] Lower default weight_decay in RL config from 0.01 to 0.001 (#5747) In full FT, AdamW weight decay shrinks the parameter directly so the implicit prior is W -> 0. In LoRA the trained parameters are A and B while the effective weight is W = W_init + (alpha/r) * B @ A; decaying A and B separately drives BA -> 0, hence W -> W_init rather than 0. The previous default of 0.01 inherited from full-FT recipes adds a measurable pull on the merged adapter back toward the base model over a few thousand steps. 0.001 keeps a small Frobenius-norm prior on ||A||^2 + ||B||^2 for numerical stability without meaningfully biasing the merged weight toward init, and aligns with the value used across the unsloth notebook templates. --- unsloth/models/rl.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 7b7c3ac1a4..c82c8364b3 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1312,7 +1312,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "logging_nan_inf_filter": False, "per_device_train_batch_size": 4, "gradient_accumulation_steps": 2, - "weight_decay": 0.01, + # LoRA decays A and B toward 0 so effective W = W_init + (alpha/r) * B @ A is pulled toward W_init, not 0 as in full FT. + # 0.001 keeps a small Frobenius prior |A|_F^2 + |B|_F^2 without measurably dragging the merged adapter back to base. + "weight_decay": 0.001, "seed": 3407, "optim": "adamw_8bit", "learning_rate": 5e-05, From f7f540a58b853d262ec7ba90c9c0af5e742cc696 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 05:00:08 -0700 Subject: [PATCH 03/43] Studio: strip orphan tool_call XML leaking into visible content (#5735) * Studio: strip orphan tool_call XML from streamed visible content The speculative-buffer state machine in `studio/backend/core/inference/llama_cpp.py` can slice a tool_call XML block between the silent DRAINING path and the user-visible content_accum, depending on when in the model's emission the BUFFERING -> STREAMING -> DRAINING transitions fire. Three leak shapes were observed in a 2026-05-22 sweep of 900 Qwen3.5 / Qwen3.6 GGUF runs: Pre-fix XML leak rate: 20/900 (2.22%), concentrated 6.7% on the larger Q8 / MTP configs: Qwen3.6-35B-A3B Q8_0 4/60 (6.7%) Qwen3.6-35B-A3B-MTP Q4 4/60 (6.7%) Qwen3.5-35B-A3B Q8_0 3/60 (5.0%) Qwen3.6-27B Q8_0 3/60 (5.0%) The existing `_TOOL_XML_RE` only matched well-formed `...` and `` pairs, so unterminated openings (close was DRAINED) and orphan closes (opening was DRAINED) survived the strip and reached the user. Fix relaxes the regex to also strip: 1. Orphan opening up to end-of-string: `(?:|\Z)` 2. Orphan closing tag: bare `` / `` Verified on the full sweep: 20/900 -> 0/900 (100% of detected leaks eliminated). 16 unit tests in `test_tool_xml_strip.py` pin all three leak shapes plus the well-formed cases, plus parametrised checks on the 5 actual real-world leak samples from the sweep data. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: strip tail-only orphan + tighten regex The 2026-05-22 gdpval sweep surfaced a 4th XML-leak shape not caught by the earlier regex: a bare `\n\n` at end-of-buffer (7 of 192 trials, all Qwen3.5-27B + a few Qwen3.6-27B). The model emits the full `...content... ` envelope, the speculative buffer DRAINS the opening tags as intended, but EOS (max_tokens cutoff) truncates the outer `` close, leaving just `` as the visible tail. We strip this ONLY when end-anchored (`\s*\Z`) so legitimate mid-text uses (user code samples, documentation discussing the Qwen tool-call XML shape) survive. Verified on the 192-trial gdpval corpus: before=7, after=0. While at it, fold the five top-level alternations into three by sharing tag-name and prefix subgroups: ... + ... + --> <(?:tool_call|function=\w+)>... | --> Semantically identical (verified by replay over the 192-trial corpus + adversarial inputs, 0 diffs) and 1.34x faster on real workloads. Backtracking-safety pinned by two new perf guards (256KB '<' spam, 1000x orphan opens). Tests: 16 -> 28 (6 new functional + 4 well-formed-vs-orphan + 2 perf guards). * Tighten comments in XML-strip regex and tests Code says what it does; comments were repeating it. Strip the verbose explanations down to the WHY-only bits (engine quirk, tail-anchor rationale, real-world source of each test sample). No code changes. inference.py: 21 -> 12 lines around _TOOL_XML_RE test_tool_xml_strip.py: 343 -> 259 lines (-84) Tests: 28/28 still pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/inference.py | 12 +- studio/backend/tests/test_tool_xml_strip.py | 263 ++++++++++++++++++++ 2 files changed, 273 insertions(+), 2 deletions(-) create mode 100644 studio/backend/tests/test_tool_xml_strip.py diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 02270ab405..bf92055929 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -427,9 +427,17 @@ _TOOL_ACTION_NUDGE = ( " Do NOT output code blocks -- use the python tool instead." ) -# Regex for stripping leaked tool-call XML from assistant messages/stream +# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py +# split across the visible/DRAIN boundary. Four leak shapes: +# 1. well-formed `...` / `...` +# 2. orphan opening to EOF (close was DRAINED) +# 3. bare orphan close (open was DRAINED) +# 4. tail-only `` (outer close truncated by EOS); anchored to +# `\Z` so mid-text `` in user code samples survives. _TOOL_XML_RE = _re.compile( - r".*?|.*?", + r"<(?:tool_call|function=\w+)>.*?(?:|\Z)" + r"|" + r"|\s*\Z", _re.DOTALL, ) logger = get_logger(__name__) diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py new file mode 100644 index 0000000000..8b90a46d5a --- /dev/null +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -0,0 +1,263 @@ +# 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 `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call +XML that leaks past the speculative buffer in core/inference/llama_cpp.py +when the open/close pair is split across the visible/DRAIN boundary. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Extract the regex from source (routes module needs heavy stubbing to import). +import re as _re + +_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() +_m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) +assert _m, "could not extract _TOOL_XML_RE source" +_ns = {"_re": _re} +exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) +_TOOL_XML_RE = _ns["_TOOL_XML_RE"] + + +# ── Well-formed pairs ───────────────────────────────────────────── + + +def test_strips_well_formed_tool_call(): + text = ( + "Let me search.\n" + "\n" + "\n" + "\nBillboard 2015\n\n" + "\n" + "\n" + "Here are the songs:" + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "" not in cleaned + assert "" not in cleaned + assert "Here are the songs:" in cleaned, "non-XML content must survive" + assert "Let me search." in cleaned + + +def test_strips_function_only_well_formed(): + text = "Setup.\n\n\nprint(1)\n\n\nDone." + cleaned = _TOOL_XML_RE.sub("", text) + assert "" + "\n" + "\n" + "\nBillboard 2015\n\n" + "" not in cleaned + assert "\n\nprint(1)\n" + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "") + assert "" not in cleaned + assert "Search starting." in cleaned + + +def test_strips_multiple_orphans(): + text = ( + "First call:\n\n\n\nx=1\n" + "Second call:\n\n\nhi\n" + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "" not in cleaned + assert "" not in cleaned + # Mid-string intentionally preserved (see preserve test). + + +# ── Tail-only (PR #5735 follow-up) ─────────────────── + + +def test_strips_tail_only_parameter_orphan(): + # Outer truncated by EOS, inner DRAINED. + cleaned = _TOOL_XML_RE.sub("", "and the text is not readable.\n\n\n") + assert "" not in cleaned + assert "and the text is not readable." in cleaned + + +def test_strips_tail_only_parameter_orphan_single_newline(): + cleaned = _TOOL_XML_RE.sub("", "Global Economic Prospects\n\n") + assert "" not in cleaned + assert "Global Economic Prospects" in cleaned + + +def test_strips_tail_only_parameter_orphan_no_trailing_ws(): + cleaned = _TOOL_XML_RE.sub("", "Final answer.") + assert "" not in cleaned + assert "Final answer." in cleaned + + +def test_preserves_mid_string_parameter_in_code_sample(): + # Tail-anchor on `` is required so doc/example prose survives. + text = ( + "Here is the Qwen tool-call format:\n" + "```xml\n" + "value\n" + "```\n" + "Note the closing sits inside ." + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "Note the closing sits inside" in cleaned + + +def test_strips_well_formed_then_orphan(): + text = ( + "Round one:\n\n\n\n1\n" + "\n\n\n" + "Now round two:\n\n\n\n" + "what is X\n\n" not in cleaned + assert "\n\n\n"Billboard Hot 100" "2015" "weekly" "chart" "position" "3"\n\n\n\n\n"peaked at number 3" Billboard Hot 100 2015 list\n\n\n\n\n"List of Billboard Hot 100 top-ten singles in 2015" wikipedia\n\n\n\nThe user wants me to list and categorize all songs that charted #3 on the Billboard Hot 100 in 2015. I have been trying to get this data", + # Qwen3.6-35B-A3B Q8_0 billboard s21 -- orphan close + "parse it more carefully.\n\n\nThe user wants a list of songs that charted #3 on the Billboard Hot 100 in 2015, categorized.", +] + + +@pytest.mark.parametrize( + "leak", REAL_LEAKS, ids = [f"sweep_sample_{i}" for i in range(len(REAL_LEAKS))] +) +def test_real_world_sweep_leaks_get_stripped(leak): + cleaned = _TOOL_XML_RE.sub("", leak) + assert "" not in cleaned, f"leak survived: {cleaned!r}" + assert " from gdpval sweep ────────── + + +# All end-anchored: outer truncated by EOS, +# inner open DRAINED, leaving bare tail. +GDPVAL_PARAMETER_LEAKS = [ + # Qwen3.5-27B Q8_0 / worldbank s00 + "the page contains image data and the text is not readable.\n\n\n", + # Qwen3.5-27B Q8_0 / worldbank s42 (preceded by mojibake) + "...some mojibake content here...\n\n\n", + # Qwen3.5-27B UD-Q4_K_XL / coppa s07 + "blocked, while others may still be in effect. The law is currently under further review by the Ninth Circuit.\n\n\n", + # Qwen3.5-27B UD-Q4_K_XL / police_training s00 + "comprehensive training report\n\n\n", + # Qwen3.5-27B UD-Q4_K_XL / worldbank s00 + "Global Economic Prospects\nJune 2025\nGlobal Economic Prospects\n\n", + # Qwen3.6-27B Q8_0 / overpass s07 + "Let me create a comprehensive query and instructions document.\n\n\n", +] + + +@pytest.mark.parametrize( + "leak", + GDPVAL_PARAMETER_LEAKS, + ids = [f"gdpval_param_orphan_{i}" for i in range(len(GDPVAL_PARAMETER_LEAKS))], +) +def test_gdpval_parameter_orphans_get_stripped(leak): + cleaned = _TOOL_XML_RE.sub("", leak) + assert "" not in cleaned, f"leak survived: {cleaned!r}" + + +# ── Backtracking guards ────────────────────────────────────────── + + +def test_no_catastrophic_backtracking_on_open_bracket_spam(): + # 256KB of '<' must fail fast (literal mismatch char 2), not backtrack. + import time + + adv = "<" * (1024 * 256) + "X" + t0 = time.perf_counter() + _TOOL_XML_RE.sub("", adv) + elapsed = time.perf_counter() - t0 + assert elapsed < 0.5, f"regex took {elapsed*1000:.0f}ms on 256KB '<' spam" + + +def test_no_catastrophic_backtracking_on_orphan_opening_spam(): + # 1000 unclosed openings: first alt must consume them all greedily. + import time + + adv = "X" * 1000 + t0 = time.perf_counter() + cleaned = _TOOL_XML_RE.sub("", adv) + elapsed = time.perf_counter() - t0 + assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens" + assert "" not in cleaned From 9c5d751c667c874b519f5dc031320b01b709bf4b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 07:11:21 -0700 Subject: [PATCH 04/43] Fixes --- unsloth/models/_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 3b67c0f487..b940fdf35a 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.6" +__version__ = "2026.5.7" __all__ = [ "SUPPORTS_BFLOAT16", From eeb49d54b8d801a1ce922fa15f778e2bad205db0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 07:22:49 -0700 Subject: [PATCH 05/43] Bump install.sh / install.ps1 pin to unsloth>=2026.5.7 (#5753) PyPI release unsloth 2026.5.7 is now live. Bumps the pinned floor in install.sh and install.ps1 from unsloth>=2026.5.6 to unsloth>=2026.5.7 so fresh installs resolve to the new wheel. Tagged on main as v0.1.416-beta. --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 07512de720..3911236d87 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1300,7 +1300,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.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.7" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1314,7 +1314,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.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.7" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1352,7 +1352,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.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.7" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1364,7 +1364,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.7" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1392,7 +1392,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.6" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.7" --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 1e19e951de..cc92fd52c2 100755 --- a/install.sh +++ b/install.sh @@ -1865,7 +1865,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.6" unsloth-zoo + "unsloth>=2026.5.7" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -1878,7 +1878,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.6" unsloth-zoo + "unsloth>=2026.5.7" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2046,7 +2046,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.6" unsloth-zoo + "unsloth>=2026.5.7" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2064,7 +2064,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.6" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.7" 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..." @@ -2096,7 +2096,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.6" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.7" --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 748aa1c482cb4c855314c4dda3119a6ce04010a8 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Mon, 25 May 2026 19:04:07 +0800 Subject: [PATCH 06/43] fix: repair mlx studio base export save_method (#5727) --- studio/backend/core/export/export.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 4ab95d896f..7cabd382eb 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -475,6 +475,7 @@ class ExportBackend: self.current_model.save_pretrained_merged( save_directory, self.current_tokenizer, + save_method = "merged_16bit", ) else: self.current_model.save_pretrained(save_directory) @@ -510,6 +511,7 @@ class ExportBackend: self.current_model.save_pretrained_merged( tmp_dir, self.current_tokenizer, + save_method = "merged_16bit", ) self.current_model.push_to_hub_merged( repo_id, From af6504f900fe611a056e66eec6ab74976eab7f34 Mon Sep 17 00:00:00 2001 From: Ricardo-M-L <69202550+Ricardo-M-L@users.noreply.github.com> Date: Mon, 25 May 2026 21:19:01 +0800 Subject: [PATCH 07/43] fix(chat_templates): check find() return value before slicing on placeholders (#5763) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(chat_templates): check find() return value before slicing on placeholders Two places in `construct_chat_template()` use `str.find()` for sentinel placeholders (`{INPUT}` / `{OUTPUT}`) without checking the -1 return: 1. The `except:` fallback (around line 2464) computes `chat_template[chat_template.find("{OUTPUT}") + len("{OUTPUT}"):]`. If the template has no `{OUTPUT}` marker, `find()` returns -1 and the slice starts at offset 7 (`-1 + len("{OUTPUT}")`), producing garbage that's then `re.escape`-d and fed back into the template-recovery regex. The user sees a confusing `IndexError` on `response_part = response_part[0]` instead of the real problem. 2. The final trim before returning (`input_part[:input_part.find("{INPUT}")]` and the matching `{OUTPUT}` line) silently drops the last character when the placeholder is missing — `find()` returns -1, and `[:-1]` slices everything except the last character, returning a corrupted template prefix to the caller. Replace both with an explicit `-1` check that raises a clear `RuntimeError` naming the missing placeholder, matching the existing guard pattern from #4923 (`try_fix_tokenizer`). Co-Authored-By: Claude Opus 4.7 * fix(chat_templates): also guard {INPUT} and fallback regex/separator paths Builds on the {OUTPUT} / final-trim guards in this branch by closing the three remaining ways the except-block fallback in construct_chat_template() can still raise a confusing IndexError or AttributeError on malformed templates: 1. Validate both {INPUT} and {OUTPUT} before deriving `ending`. The regex two lines later (`{INPUT} + ending + ...`) still produced an empty list and crashed on `response_part[0]` if {INPUT} was missing. 2. Guard the regex no-match case. Some templates contain both placeholders but not in a recoverable two-example shape, in which case `re.findall` returns an empty list and `[0]` raises. 3. Initialize `found = None` before the separator-search loop and raise if the loop never sets it. Previously, if the first iteration's `re.finditer` was empty the loop broke without binding `found`, and `found.group(1)` raised AttributeError on the stale int left over from the outer rfind loop. Rephrase the final-trim error messages from internal variable names ("input_part") to user-facing wording ("instruction section") and include a bounded (200-char) excerpt of the offending content so the error is debuggable without being unbounded. Add tests/python/test_construct_chat_template_validation.py covering each failure mode with a fake tokenizer (no HF_TOKEN, no model download, CPU-only). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Claude Opus 4.7 Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- ...test_construct_chat_template_validation.py | 77 +++++++++++++++++++ unsloth/chat_templates.py | 41 +++++++++- 2 files changed, 115 insertions(+), 3 deletions(-) create mode 100644 tests/python/test_construct_chat_template_validation.py diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py new file mode 100644 index 0000000000..9ab68639c4 --- /dev/null +++ b/tests/python/test_construct_chat_template_validation.py @@ -0,0 +1,77 @@ +"""Negative-path validation tests for unsloth.chat_templates.construct_chat_template. + +Regression coverage for the str.find() / regex no-match guards added in +PR #5763 follow-up: missing placeholders or unrecoverable two-example +structures must raise RuntimeError with a clear message, not IndexError +or AttributeError, and must never silently drop the last character via +s[:-1]. + +Uses a minimal fake tokenizer so the cases run on CPU-only CI without +HF_TOKEN and without downloading a gated model. The validation paths +exercised here fail before construct_chat_template reaches any heavy +tokenizer interaction, so the stub stays small. +""" + +import pytest + +from unsloth.chat_templates import construct_chat_template + + +class _FakeTokenizer: + """Minimum surface construct_chat_template touches before the + validation guards fire.""" + + name_or_path = "fake/tokenizer" + eos_token = "" + + def get_vocab(self): + return {"": 0} + + +@pytest.mark.parametrize( + "template, expected_in_message", + [ + ("only {INPUT} here, no output marker", "{OUTPUT}"), + ("only {OUTPUT} here, no input marker", "{INPUT}"), + ("neither sentinel here, just literal text", "{INPUT}"), + ("neither sentinel here, just literal text", "{OUTPUT}"), + ], +) +def test_missing_placeholder_in_chat_template_raises(template, expected_in_message): + with pytest.raises(RuntimeError) as exc_info: + construct_chat_template( + tokenizer = _FakeTokenizer(), + chat_template = template, + extra_eos_tokens = [""], + ) + assert expected_in_message in str(exc_info.value) + + +def test_single_pair_template_raises_clear_error_not_attribute_error(): + """One {INPUT}/{OUTPUT} pair (rather than the required two) used to + crash with AttributeError on `found.group(1)` after the for-loop + broke without setting `found`. Must raise RuntimeError now.""" + template = "user: {INPUT}\nassistant: {OUTPUT}\n" + with pytest.raises(RuntimeError): + construct_chat_template( + tokenizer = _FakeTokenizer(), + chat_template = template, + extra_eos_tokens = [""], + ) + + +def test_error_message_excerpt_is_bounded(): + """Error messages must include a bounded excerpt of the offending + template, not dump arbitrarily large content into the traceback.""" + huge = ("garbage " * 5000) + "{INPUT}" # ~40 KB, missing {OUTPUT} + with pytest.raises(RuntimeError) as exc_info: + construct_chat_template( + tokenizer = _FakeTokenizer(), + chat_template = huge, + extra_eos_tokens = [""], + ) + msg = str(exc_info.value) + # Excerpt is repr-quoted and capped; total message should stay well + # under the template length. + assert len(msg) < 1000 + assert "{OUTPUT}" in msg diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index e8a34cbc60..956fcb2392 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2461,17 +2461,40 @@ extra_eos_tokens = None, f"{left_changed}" ) except: - ending = chat_template[chat_template.find("{OUTPUT}") + len("{OUTPUT}"):] + output_pos = chat_template.find("{OUTPUT}") + input_pos = chat_template.find("{INPUT}") + if output_pos == -1 or input_pos == -1: + missing = [] + if input_pos == -1: missing.append("{INPUT}") + if output_pos == -1: missing.append("{OUTPUT}") + raise RuntimeError( + f"Unsloth: chat_template must contain {' and '.join(missing)} " + f"placeholder(s). Got: {chat_template[:200]!r}" + ) + ending = chat_template[output_pos + len("{OUTPUT}"):] ending = re.escape(ending) find_text = "{INPUT}" + ending + "(.+?{OUTPUT}" + ending + ")" response_part = re.findall(find_text, chat_template, flags = re.DOTALL | re.MULTILINE) + if len(response_part) == 0: + raise RuntimeError( + "Unsloth: Could not recover a two-example structure from chat_template. " + "Provide exactly two {INPUT}/{OUTPUT} pairs (and optionally {SYSTEM}). " + f"Got: {chat_template[:200]!r}" + ) response_part = response_part[0] + found = None for j in range(1, len(response_part)): try_find = re.escape(response_part[:j]) try: found = next(re.finditer("(" + try_find + ").+?\\{INPUT\\}", chat_template, flags = re.DOTALL | re.MULTILINE)) except: break + if found is None: + raise RuntimeError( + "Unsloth: Could not locate a separator between examples in chat_template. " + "Provide exactly two {INPUT}/{OUTPUT} pairs (and optionally {SYSTEM}). " + f"Got: {chat_template[:200]!r}" + ) separator = found.group(1) response_start = chat_template.find(response_part) @@ -2607,8 +2630,20 @@ extra_eos_tokens = None, jinja_template = "{{ bos_token }}" + jinja_template # Get instruction and output parts for train_on_inputs = False - input_part = input_part [:input_part .find("{INPUT}")] - output_part = output_part[:output_part.find("{OUTPUT}")] + input_idx = input_part .find("{INPUT}") + output_idx = output_part.find("{OUTPUT}") + if input_idx == -1: + raise RuntimeError( + f"Unsloth: The instruction section of the template must contain the " + f"'{{INPUT}}' placeholder. Section: {input_part[:200]!r}" + ) + if output_idx == -1: + raise RuntimeError( + f"Unsloth: The response section of the template must contain the " + f"'{{OUTPUT}}' placeholder. Section: {output_part[:200]!r}" + ) + input_part = input_part [:input_idx ] + output_part = output_part[:output_idx] return modelfile, jinja_template, input_part, output_part From b73480e55467f48d628a4a91a21e045c95489740 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 22:21:43 -0700 Subject: [PATCH 08/43] [pre-commit.ci] pre-commit autoupdate (#5773) 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.13 → v0.15.14](https://github.com/astral-sh/ruff-pre-commit/compare/v0.15.13...v0.15.14) 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 9d80fe6ff5..1919fac9c5 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.13 + rev: v0.15.14 hooks: - id: ruff args: From 034ff512e7269e7c3cedfd0f8c388065a2384537 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 23:36:51 -0700 Subject: [PATCH 09/43] Studio: stop seeded admin to cross-origin callers (#5739) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: stop leaking seeded admin pw to cross-origin callers The "/" SPA fallback serves index.html with an inline ``window.__UNSLOTH_BOOTSTRAP__`` script containing the seeded admin password while a password change is pending. Default web mode runs ``CORSMiddleware`` with ``allow_origins=["*"]`` + ``allow_credentials= True``, which reflects an attacker-controlled ``Origin`` back on every request and sets Access-Control-Allow-Credentials true. The combination let any cross-origin page ``fetch('/')`` with credentials and read the bootstrap admin password out of the HTML body. The API smoke ``CORS: GET / leaks bootstrap pw to cross-origin caller`` audit already tracked this (tests/studio/studio_api_smoke.py:224) but did not gate CI. Gate ``_inject_bootstrap`` on a same-origin check: legitimate top-level navigations omit ``Origin`` on most engines, so the absence of the header is treated as same-origin; when the header IS present and does not match ``request.url.scheme://request.url.netloc`` exactly, we now skip injecting the bootstrap tag. ``Vary: Origin`` is added so an intermediary cache cannot serve a same-origin response (with bootstrap) to a later cross-origin caller (and vice versa). Coverage: ``test_index_bootstrap_origin.py`` exercises the helper with missing / matching / evil / scheme-mismatch / port-mismatch origins. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten comments on bootstrap cross-origin helper * Studio: canonicalise Origin before same-origin gate A plain string-compare between the Origin header and request.url.netloc misclassifies legitimate same-origin requests as cross-origin in three scenarios: - Browser strips the default port from Origin (https://example.com) but Starlette's netloc keeps it (example.com:443). Per RFC 6454 the default port is dropped on the wire, so the strings will not match even though the requests share an origin. - Host case differs (Origin: http://Example.com vs netloc: example.com). Per RFC 3986 host comparison is case-insensitive. - Scheme case differs (HTTP:// vs http://). Per RFC 3986 the scheme is also case-insensitive. These are usability degradations rather than security gaps (legitimate user denied the bootstrap injection, no attacker gain), but worth shipping so non-default Studio deployments keep the change-password auto-fill. Adds _canonical_origin(scheme, netloc) -> (scheme, host, port) and compares the canonical tuples. Default-port lookup covers http/https/ws/wss; userinfo (user:pass@) is stripped per RFC 3986 since Origin never carries credentials. Origin: "null" (sandboxed iframes, file:// pages) and unparseable values collapse to cross- origin so the bootstrap pw is never leaked through those paths either. Tests: 14 cases (was 5). Covers the original same/missing/evil/ scheme/port matrix plus default-port stripping in both directions, host + scheme case folding, Origin: null, garbage values, and userinfo-in-netloc. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix IPv6 netloc parsing for PR #5739 The canonical-origin helper used ``netloc.partition(":")`` which mis-parses bracketed IPv6 hosts (``[::1]:8902`` -> host=``[``, port-str=``:1]:8902``). The int() then raises and the canonicaliser returns None, so every IPv6 same-origin request is misclassified as cross-origin and Studio refuses to inject the bootstrap pw on a legitimate top-level nav when launched with ``unsloth studio -H ::1``. Bracket-aware split per RFC 3986 §3.2.2, plus extra regression tests for IPv6, opaque (data:/blob:/file:), comma-joined multi-Origin and localhost-vs-127 cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard urlparse ValueError in same-origin gate urlparse raises ValueError on malformed bracketed Origin values (unclosed [, invalid IPv6 hex, text after ]) and on a few NFKC edge cases since Py 3.8. Without a guard, a request carrying Origin: http://[malformed surfaced as HTTP 500 from the SPA handler rather than being treated as cross-origin per the docstring's safer-default rule. Wrap both urlparse calls in try/except ValueError and return False on parse failure. Also distinguish a missing Origin header (top-level same-document GET, treat as same-origin) from an explicit empty string (not a valid serialised origin per RFC 6454 §6.1, treat as cross-origin). Four new regression tests pinned down by the PR audit: malformed IPv6 bracket, invalid IPv6 hex, bracket with trailing garbage, and the empty Origin header. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shorten origin-gate comments for PR #5739 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/main.py | 109 +++++++++- .../tests/test_index_bootstrap_origin.py | 144 +++++++++++++ .../test_index_bootstrap_origin_extra.py | 196 ++++++++++++++++++ 3 files changed, 438 insertions(+), 11 deletions(-) create mode 100644 studio/backend/tests/test_index_bootstrap_origin.py create mode 100644 studio/backend/tests/test_index_bootstrap_origin_extra.py diff --git a/studio/backend/main.py b/studio/backend/main.py index 004ae404cd..689241b915 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -49,6 +49,8 @@ import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version +from typing import Optional +from urllib.parse import urlparse _STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$") @@ -715,10 +717,8 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes: def _inject_bootstrap(html_bytes: bytes, app: FastAPI): """Inject bootstrap credentials when password change is pending. - - Returns ``(html_bytes, script_nonce_or_None)``. Callers must forward - the nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so the inline script is - not blocked by CSP. + Returns ``(html_bytes, script_nonce_or_None)``; callers forward the + nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so CSP allows the inline script. """ import json as _json import secrets as _secrets @@ -743,6 +743,86 @@ def _inject_bootstrap(html_bytes: bytes, app: FastAPI): return html.encode("utf-8"), nonce +_DEFAULT_PORTS = {"http": 80, "https": 443, "ws": 80, "wss": 443} + + +def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]]: + """Canonicalise an Origin to ``(scheme, host, port)`` for equality. + Browsers strip default ports (RFC 6454 sec 6.1) and scheme/host are + case-insensitive (RFC 3986), so bare string compare misclassifies + same-origin requests as cross-origin. Returns ``None`` on unparseable + input so callers fall to the safer cross-origin default. + """ + scheme = (scheme or "").strip().lower() + if not scheme or not netloc: + return None + # Strip userinfo (RFC 3986); Origin never carries credentials. + if "@" in netloc: + netloc = netloc.rsplit("@", 1)[1] + # IPv6 hosts use brackets (RFC 3986 sec 3.2.2): ``[::1]:8902``. Bare + # ``partition(":")`` mis-parses these and breaks ``unsloth studio -H ::1``. + if netloc.startswith("["): + close = netloc.find("]") + if close == -1: + return None + host = netloc[1:close] + rest = netloc[close + 1 :] + if rest.startswith(":"): + port_str = rest[1:] + elif rest == "": + port_str = "" + else: + return None + else: + host, _, port_str = netloc.partition(":") + host = host.strip().lower() + if not host: + return None + if port_str: + try: + port = int(port_str) + except ValueError: + return None + else: + port = _DEFAULT_PORTS.get(scheme, 0) + return (scheme, host, port) + + +def _is_same_origin_request(request: Request) -> bool: + """True when Origin is missing or matches request's scheme://host:port. + Top-level same-document GETs omit Origin, so missing counts as same-origin. + Callers must also emit ``Vary: Origin``. Both sides are canonicalised via + :func:`_canonical_origin` so default-port stripping and scheme/host case + do not misclassify same-origin requests as cross-origin. + """ + origin = request.headers.get("origin") + if origin is None: + # Missing header: top-level same-document GETs omit Origin. + return True + # Empty string is not a valid serialised origin (RFC 6454 sec 6.1). + if not origin: + return False + # "null" token (sandboxed iframes, file:// pages) is never same-origin. + if origin == "null": + return False + # ``urlparse`` raises ``ValueError`` on malformed IPv6 brackets; swallow + # so a garbage Origin doesn't 500 the SPA handler. + try: + parsed = urlparse(origin) + except ValueError: + return False + origin_canon = _canonical_origin(parsed.scheme, parsed.netloc) + if origin_canon is None: + return False + try: + self_canon = _canonical_origin(request.url.scheme, request.url.netloc) + except ValueError: + return False + if self_canon is None: + return False + return origin_canon == self_canon + + def setup_frontend(app: FastAPI, build_path: Path): """Mount frontend static files (optional)""" if not build_path.exists(): @@ -753,11 +833,18 @@ def setup_frontend(app: FastAPI, build_path: Path): if assets_dir.exists(): app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") - def _build_index_response() -> Response: + def _build_index_response(request: Request) -> Response: content = (build_path / "index.html").read_bytes() content = _strip_crossorigin(content) - content, nonce = _inject_bootstrap(content, app) - headers = {"Cache-Control": "no-cache, no-store, must-revalidate"} + # Bootstrap pw is same-origin only; Vary: Origin keeps caches honest. + if _is_same_origin_request(request): + content, nonce = _inject_bootstrap(content, app) + else: + nonce = None + headers = { + "Cache-Control": "no-cache, no-store, must-revalidate", + "Vary": "Origin", + } if nonce: headers[_CSP_SCRIPT_NONCE_HEADER] = nonce return Response( @@ -767,11 +854,11 @@ def setup_frontend(app: FastAPI, build_path: Path): ) @app.get("/") - async def serve_root(): - return _build_index_response() + async def serve_root(request: Request): + return _build_index_response(request) @app.get("/{full_path:path}") - async def serve_frontend(full_path: str): + async def serve_frontend(request: Request, full_path: str): if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")): return {"error": "API endpoint not found"} @@ -785,6 +872,6 @@ def setup_frontend(app: FastAPI, build_path: Path): return FileResponse(file_path) # Serve index.html as bytes — avoids Content-Length mismatch - return _build_index_response() + return _build_index_response(request) return True diff --git a/studio/backend/tests/test_index_bootstrap_origin.py b/studio/backend/tests/test_index_bootstrap_origin.py new file mode 100644 index 0000000000..89f7613ee4 --- /dev/null +++ b/studio/backend/tests/test_index_bootstrap_origin.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression coverage for the bootstrap-pw cross-origin leak (PR 5739). +``_is_same_origin_request`` gates ``_inject_bootstrap`` so the seeded +admin password only ships to same-origin callers. +""" + +import os +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + + +def _build_request(host: str, origin: str | None, scheme: str = "http") -> MagicMock: + request = MagicMock() + request.url.scheme = scheme + request.url.netloc = host + request.headers = {"origin": origin} if origin is not None else {} + return request + + +def test_is_same_origin_request_missing_origin_is_same_origin(monkeypatch): + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = None) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_matching_origin_is_same_origin(): + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "http://127.0.0.1:8888") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_evil_origin_is_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "https://evil.example") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_scheme_mismatch_is_cross_origin(): + # https origin against an http listener is not same-origin. + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "https://127.0.0.1:8888") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_port_mismatch_is_cross_origin(): + # Same host different port is not same-origin per the web platform. + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "http://127.0.0.1:5173") + assert _is_same_origin_request(req) is False + + +# ── Canonicalisation: default-port stripping + case folding ───────── + + +def test_is_same_origin_request_https_default_port_stripped_on_origin(): + """RFC 6454 strips default ports on Origin; Starlette's netloc may still + carry ``:443``. Canonicalise both sides so this stays same-origin. + """ + from main import _is_same_origin_request + + req = _build_request( + "example.com:443", origin = "https://example.com", scheme = "https" + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_http_default_port_stripped_on_origin(): + from main import _is_same_origin_request + + req = _build_request("example.com:80", origin = "http://example.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_default_port_present_on_origin(): + """Mirror case: Origin carries the default port, netloc doesn't. Same-origin.""" + from main import _is_same_origin_request + + req = _build_request( + "example.com", origin = "https://example.com:443", scheme = "https" + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_host_case_insensitive(): + """Host portion is case-insensitive per RFC 3986.""" + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "http://EXAMPLE.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_scheme_case_insensitive(): + """Scheme portion is case-insensitive per RFC 3986.""" + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "HTTP://example.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_null_origin_is_cross_origin(): + """Sandboxed iframes / file:// pages send ``Origin: null``; cross-origin.""" + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "null") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_unparseable_origin_is_cross_origin(): + """Garbage values without a host fall to cross-origin; a malformed header + must not leak the bootstrap. + """ + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "not-a-url") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_userinfo_in_netloc_ignored(): + """``user:pass@host:port`` netlocs (RFC 3986) must compare equal to the + credentials-less Origin. + """ + from main import _is_same_origin_request + + req = _build_request("user:pass@example.com:80", origin = "http://example.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_explicit_non_default_port_still_mismatch(): + """Canonicalisation does NOT collapse non-default ports to default.""" + from main import _is_same_origin_request + + req = _build_request( + "example.com", origin = "https://example.com:9999", scheme = "https" + ) + assert _is_same_origin_request(req) is False diff --git a/studio/backend/tests/test_index_bootstrap_origin_extra.py b/studio/backend/tests/test_index_bootstrap_origin_extra.py new file mode 100644 index 0000000000..aea6b36a96 --- /dev/null +++ b/studio/backend/tests/test_index_bootstrap_origin_extra.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Extra edge-case coverage for the bootstrap-pw cross-origin gate. +Companion to ``test_index_bootstrap_origin.py``: IPv6 netlocs, opaque +origins (``data:``, ``blob:``), comma-joined multi-Origin headers, and +the ``localhost`` vs ``127.0.0.1`` distinct-origin rule. +""" + +from unittest.mock import MagicMock + + +def _build_request(host: str, origin, scheme: str = "http") -> MagicMock: + request = MagicMock() + request.url.scheme = scheme + request.url.netloc = host + request.headers = {"origin": origin} if origin is not None else {} + return request + + +# ── IPv6 ──────────────────────────────────────────────────────────── + + +def test_is_same_origin_request_ipv6_loopback_same_origin(): + """Studio supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare + ``partition(":")`` mis-parses the bracketed form and would refuse the + bootstrap on legitimate same-origin nav. + """ + from main import _is_same_origin_request + + req = _build_request("[::1]:8902", origin = "http://[::1]:8902") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_full_address_same_origin(): + from main import _is_same_origin_request + + req = _build_request( + "[2001:db8::1]:8443", + origin = "https://[2001:db8::1]:8443", + scheme = "https", + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_default_port_stripped(): + """Browser drops :80 on ``http://[::1]``.""" + from main import _is_same_origin_request + + req = _build_request("[::1]:80", origin = "http://[::1]") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_case_insensitive(): + """Hex digits in IPv6 are case-insensitive per RFC 5952.""" + from main import _is_same_origin_request + + req = _build_request( + "[2001:DB8::1]:8443", + origin = "https://[2001:db8::1]:8443", + scheme = "https", + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_different_host_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("[::1]:8902", origin = "http://[2001:db8::1]:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_ipv6_port_mismatch_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("[::1]:8902", origin = "http://[::1]:9999") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_ipv6_userinfo_stripped(): + from main import _is_same_origin_request + + req = _build_request("user:pass@[::1]:8902", origin = "http://[::1]:8902") + assert _is_same_origin_request(req) is True + + +# ── Opaque origins (data:, blob:) ─────────────────────────────────── + + +def test_is_same_origin_request_data_url_origin_is_cross_origin(): + """``data:`` URLs are opaque origins (HTML living standard); no host, + never same-origin. + """ + from main import _is_same_origin_request + + req = _build_request( + "127.0.0.1:8902", origin = "data:text/html," + ) + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_blob_url_origin_is_cross_origin(): + """``blob:`` URLs carry the inner origin only in non-canonical form; the + canonical comparison rejects them. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "blob:http://127.0.0.1:8902/uuid") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_file_url_origin_is_cross_origin(): + """``file://`` pages usually send ``Origin: null``; historical engines + sent ``Origin: file://``. Neither is same-origin vs an http listener. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "file://") + assert _is_same_origin_request(req) is False + + +# ── Multi-Origin header (comma-joined by Starlette) ──────────────── + + +def test_is_same_origin_request_comma_joined_origins_cross_origin(): + """Starlette concatenates repeated headers with ``, ``; the canonical + parser can't safely split this, so it falls to cross-origin. + """ + from main import _is_same_origin_request + + req = _build_request( + "127.0.0.1:8902", + origin = "http://127.0.0.1:8902, http://evil.example", + ) + assert _is_same_origin_request(req) is False + + +# ── localhost vs 127.0.0.1 (distinct origins per web platform) ────── + + +def test_is_same_origin_request_localhost_vs_127_is_cross_origin(): + """Browsers treat ``localhost`` and ``127.0.0.1`` as distinct origins; + the canonical comparison must not DNS-collapse them. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://localhost:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_127_vs_localhost_is_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("localhost:8902", origin = "http://127.0.0.1:8902") + assert _is_same_origin_request(req) is False + + +# ── urlparse ValueError robustness ───────────────────────────────── + + +def test_is_same_origin_request_malformed_ipv6_bracket_is_cross_origin(): + """``urlparse`` raises ``ValueError('Invalid IPv6 URL')`` on unclosed + brackets (CVE-2024-11168 hardening). The gate must swallow and fall to + cross-origin rather than 500 the SPA handler. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://[malformed") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_invalid_ipv6_address_is_cross_origin(): + """Bracketed but invalid IPv6 (e.g. ``[::g]``) also raises + ``ValueError`` inside ``urlparse``.""" + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://[::g]:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_bracket_with_trailing_garbage_is_cross_origin(): + """Text after the closing bracket also raises inside ``urlparse``.""" + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://[2001:db8::1]extra:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_empty_origin_header_is_cross_origin(): + """Explicit empty ``Origin:`` is not a valid serialised origin and must + not be conflated with a missing header; cross-origin, bootstrap withheld. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "") + assert _is_same_origin_request(req) is False From cc68720385e5c3673c83077b94daad7d104a9f85 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 23:37:04 -0700 Subject: [PATCH 10/43] Studio: surface external-provider cache hits and writes in context bar (#5736) * Studio: surface external-provider cache hits and writes in context bar The Anthropic / OpenAI Responses streaming paths already emit an include_usage-style SSE chunk carrying prompt_tokens_details.cached_tokens and cache_creation_input_tokens / cache_read_input_tokens (see _build_usage_chunk in external_provider.py), but the chat-adapter only read the local llama-server timings.cache_n field. As a result, the context-usage tooltip never showed cache hits or writes for external providers, even though the backend was computing them. Read the external usage envelope as a fallback when timings.cache_n is absent, and surface Anthropic cache_creation_input_tokens as a separate "Cache writes" line in the tooltip so users can tell a cache miss from a cache hit on a turn that both reads and writes the cache. - ServerUsage gains optional prompt_tokens_details.cached_tokens, cache_creation_input_tokens, cache_read_input_tokens. - contextUsage store entry gains optional cacheWriteTokens. - ContextUsageBar gains optional cacheWrites tooltip line. - chat-page wires both fields through to the bar. * Studio: render cache stats for external providers too Reviewer round on the original PR caught three asymmetric-fix sites where the producer side surfaced external prompt-cache stats but the consumer side still gated on ggufContextLength (which is only ever set for the local llama-server runtime). Result: the entire cache-stats PR shipped invisible for Anthropic / OpenAI Responses / Gemini, which is exactly the set of providers it was added for. - chat-page.tsx: drop the ggufContextLength precondition on the ContextUsageBar mount. The bar already tracks usage; let it decide what to render based on what it knows. - context-usage-bar.tsx: make `total` optional. When absent, drop the "/ total" ratio + percentage progress bar + "approaching limit" helper, and just show per-turn counters + cache stats. Bootstrap guard tightened so an all-zero, all-undefined state still renders nothing. - runtime-provider.tsx: external-provider rehydration was rejected by the `store.ggufContextLength` check. Keep the "fits inside window" sanity check when a local context window IS known, drop it when it isn't. - message-timing.tsx: the per-message timing popover used a separate "Cache hits" code path that only read llama-server's timings.cache_n. Fall through to custom.contextUsage for external providers, and add a parallel "Cache writes" line for Anthropic cache_creation events. * Studio: tighten cache-stats comments * Scope contextUsage to active checkpoint Three follow-ups on #5736 so the relaxed external-provider render gate does not show stale token / cache stats from a different model: 1) setCheckpoint now clears contextUsage on a real checkpoint change. setActiveThreadId and clearCheckpoint already did this; the most-traveled transition path (the user switching models from the picker) leaked the prior turn's counts because they were never cleared. 2) The external-selection branch in chat-page.tsx now also clears contextUsage at the same time it nulls ggufContextLength / activeNativePathToken. Without this an in-session switch from a local model to an external provider would visibly carry the previous local turn's counters into the new provider's bar. 3) exitCompare's rehydration is now scoped: restore the saved usage only when the message's modelId matches the active checkpoint AND, for local turns where a context window is known, when the saved total fits inside that window. Without this the bar could render a stale local-model usage on top of an external provider, or an oversized usage object that exceeds the now- active window. Typecheck clean. * Plug remaining stale-contextUsage paths Follow-up to 042e0ac4 that catches four asymmetric-fix sites the checkpoint-scoping pass missed: 1) setParams now also clears contextUsage on a real checkpoint change. The local model load path in use-chat-model-runtime calls setParams(mergeBackendRecommendedInference(...)) which mutates params.checkpoint before refresh() eventually fires setCheckpoint; the intermediate window rendered the previous model's counters under the new checkpoint. 2) chat-adapter.ts setContextUsage on stream completion now gates on the captured params.checkpoint still being active. A late completion from provider A used to clobber the context bar after the user switched to provider B mid-stream. 3) chat-page.tsx exitCompare rehydration no longer accepts a saved modelId-stamped usage when the active checkpoint is empty. A user who entered compare, cleared the model, and exited compare would otherwise see the cleared model's stats reappear. 4) runtime-provider.tsx thread-load no longer restores legacy unscoped usage (no modelId) unless a local context window is known. With the relaxed external-provider render gate, old pre-PR persisted messages without a modelId stamp could attach their counts to an unrelated active provider. Also switches message-timing.tsx cache-hit fallback from || to ?? so an explicit cache_n=0 is not replaced by a stale cachedTokens. Typecheck clean. * Shorten cache-stats comments for PR #5736 --- .../assistant-ui/message-timing.tsx | 50 +++++++++-- .../src/features/chat/api/chat-adapter.ts | 29 +++++- .../frontend/src/features/chat/chat-page.tsx | 36 +++++++- .../chat/components/context-usage-bar.tsx | 89 ++++++++++++++----- .../src/features/chat/runtime-provider.tsx | 21 +++-- .../chat/stores/chat-runtime-store.ts | 16 +++- 6 files changed, 197 insertions(+), 44 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 4fb68d4b90..31f742bc5e 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -33,10 +33,24 @@ export const MessageTiming: FC<{ if (timing?.totalStreamTime === undefined) return null; - const serverTimings = ( + const custom = ( message.metadata as Record | undefined - )?.custom as { serverTimings?: Record } | undefined; - const st = serverTimings?.serverTimings; + )?.custom as + | { + serverTimings?: Record; + contextUsage?: { + cachedTokens?: number; + cacheWriteTokens?: number; + }; + } + | undefined; + const st = custom?.serverTimings; + // `??` (not `||`) so an explicit cache_n=0 isn't replaced by a stale + // contextUsage.cachedTokens from a prior turn. + const cacheHits = + st?.cache_n ?? custom?.contextUsage?.cachedTokens ?? 0; + // Anthropic-only cache-write count. + const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0; // Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op // turns, blowing the rate up to Infinity. Require >=1 token AND a @@ -122,11 +136,19 @@ export const MessageTiming: FC<{ )} - {(st?.cache_n ?? 0) > 0 && ( + {cacheHits > 0 && (
Cache hits - {formatNumber(st!.cache_n)} + {formatNumber(cacheHits)} + +
+ )} + {cacheWrites > 0 && ( +
+ Cache writes + + {formatNumber(cacheWrites)}
)} @@ -146,7 +168,7 @@ export const MessageTiming: FC<{ ) : ( <> - {/* Client-side metrics (safetensors fallback) */} + {/* Client-side metrics (safetensors + external provider fallback) */} {timing.firstTokenTime !== undefined && (
First token @@ -155,6 +177,22 @@ export const MessageTiming: FC<{
)} + {cacheHits > 0 && ( +
+ Cache hits + + {formatNumber(cacheHits)} + +
+ )} + {cacheWrites > 0 && ( +
+ Cache writes + + {formatNumber(cacheWrites)} + +
+ )}
Total diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0c557f1b01..9842e380e0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -70,6 +70,13 @@ interface ServerUsage { prompt_tokens: number; completion_tokens: number; total_tokens: number; + // External prompt-cache fields (see _build_usage_chunk in + // external_provider.py). cache_creation is Anthropic-only. + prompt_tokens_details?: { + cached_tokens?: number; + }; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; } /** Server-side timing data from llama-server's timings object. */ @@ -1881,18 +1888,31 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const finalTokPerSec = meta?.timings?.predicted_per_second; const serverPromptEvalTime = meta?.timings?.prompt_ms; - // Update context usage in store if we got valid server data + // Prefer llama-server timings; fall back to provider usage envelope. + const cachedTokens = + meta?.timings?.cache_n ?? + meta?.usage?.prompt_tokens_details?.cached_tokens ?? + meta?.usage?.cache_read_input_tokens ?? + 0; + // Anthropic-only (billed at the write premium). + const cacheWriteTokens = meta?.usage?.cache_creation_input_tokens ?? 0; + + // Gate on the captured checkpoint still being active so a late + // completion from provider A doesn't populate the bar after the + // user switched to provider B mid-stream. if ( meta?.usage && typeof meta.usage.prompt_tokens === "number" && typeof meta.usage.completion_tokens === "number" && - typeof meta.usage.total_tokens === "number" + typeof meta.usage.total_tokens === "number" && + useChatRuntimeStore.getState().params.checkpoint === params.checkpoint ) { useChatRuntimeStore.getState().setContextUsage({ promptTokens: meta.usage.prompt_tokens, completionTokens: meta.usage.completion_tokens, totalTokens: meta.usage.total_tokens, - cachedTokens: meta.timings?.cache_n ?? 0, + cachedTokens, + cacheWriteTokens, }); } @@ -1922,7 +1942,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { promptTokens: meta.usage.prompt_tokens, completionTokens: meta.usage.completion_tokens, totalTokens: meta.usage.total_tokens, - cachedTokens: meta.timings?.cache_n ?? 0, + cachedTokens, + cacheWriteTokens, modelId: params.checkpoint, } : undefined, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index ce02b0da18..85ed0f7eef 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1037,6 +1037,10 @@ export function ChatPage(): ReactElement { ggufMaxContextLength: null, ggufNativeContextLength: null, activeNativePathToken: null, + // Clear previous-model counters; the relaxed external-provider + // render gate would otherwise show stale stats until the next + // completion overwrites them. + contextUsage: null, supportsReasoning: reasoningCaps.supportsReasoning, reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn, reasoningStyle: reasoningCaps.reasoningStyle, @@ -1161,7 +1165,9 @@ export function ChatPage(): ReactElement { if (!saved) return; viewBeforeCompareRef.current = null; navigate({ to: "/chat", search: saved }); - // Restore context usage from the active thread's last assistant message. + // Restore usage from the last assistant message, but only if it + // matches the currently active checkpoint. Without this guard the + // relaxed render gate would show stale stats from another model. const threadId = saved.thread ?? useChatRuntimeStore.getState().activeThreadId; if (threadId) { @@ -1175,7 +1181,29 @@ export function ChatPage(): ReactElement { const usage = metadata?.contextUsage as ReturnType< typeof useChatRuntimeStore.getState >["contextUsage"]; - if (usage) useChatRuntimeStore.getState().setContextUsage(usage); + if (!usage) return; + const store = useChatRuntimeStore.getState(); + const activeCheckpoint = store.params.checkpoint; + const usageModelId = + (usage as { modelId?: unknown }).modelId; + // Scope by modelId when present; reject if no active checkpoint + // (model-scoped usage cannot be attributed to "nothing"). + if (typeof usageModelId === "string" && usageModelId) { + if (!activeCheckpoint || usageModelId !== activeCheckpoint) { + return; + } + } + // For local turns, also require the restored count to fit in + // the active window. Skip when unknown (external provider). + const limit = store.ggufContextLength; + if ( + typeof limit === "number" && + limit > 0 && + (usage.totalTokens ?? 0) > limit + ) { + return; + } + store.setContextUsage(usage); }) .catch((error) => { if (!isExpectedBackgroundChatStorageError(error)) { @@ -1491,11 +1519,13 @@ export function ChatPage(): ReactElement { ) : null}
- {view.mode === "single" && ggufContextLength && contextUsage ? ( + {view.mode === "single" && contextUsage ? ( = ({ used, total, cached, promptTokens, completionTokens, className }) => { - if (total <= 0) return null; +}> = ({ + used, + total, + cached, + cacheWrites, + promptTokens, + completionTokens, + className, +}) => { + const hasKnownLimit = typeof total === "number" && total > 0; + const hasUsageDetails = + promptTokens !== undefined || + completionTokens !== undefined || + (cached !== undefined && cached > 0) || + (cacheWrites !== undefined && cacheWrites > 0); - const percent = Math.min((used / total) * 100, 100); - const severity = getSeverityColor(percent); + // Nothing to show: no limit and no per-turn counters. + if (!hasKnownLimit && used <= 0 && !hasUsageDetails) return null; + + const percent = hasKnownLimit + ? Math.min((used / (total as number)) * 100, 100) + : null; + const severity = getSeverityColor(percent ?? 0); return (
-
- Context usage - - {percent.toFixed(1)}% - -
+ {hasKnownLimit && percent !== null ? ( +
+ Context usage + + {percent.toFixed(1)}% + +
+ ) : null} {promptTokens !== undefined && (
Prompt tokens @@ -98,20 +129,32 @@ export const ContextUsageBar: FC<{
)} + {cacheWrites !== undefined && cacheWrites > 0 && ( +
+ Cache writes + + {formatTokenCountFull(cacheWrites)} + +
+ )}
- Total + + {hasKnownLimit ? "Total" : "Total tokens"} + - {formatTokenCountFull(used)} / {formatTokenCountFull(total)} + {hasKnownLimit + ? `${formatTokenCountFull(used)} / ${formatTokenCountFull(total as number)}` + : formatTokenCountFull(used)}
- {percent > 85 && ( + {hasKnownLimit && percent !== null && percent > 85 ? (
Close to the context limit. Generation will stop at 100%. Increase Context Length in the chat Settings panel to keep going.
- )} + ) : null}
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index d01383b309..21be5f6e3e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -826,17 +826,24 @@ function useStudioRuntimeAdapters(): StudioRuntimeAdapters { completionTokens: number; totalTokens: number; cachedTokens: number; + cacheWriteTokens?: number; modelId?: string; } | undefined; const store = useChatRuntimeStore.getState(); - if ( - savedUsage && - store.ggufContextLength && - savedUsage.totalTokens <= store.ggufContextLength && - (!savedUsage.modelId || - savedUsage.modelId === store.params.checkpoint) - ) { + // Window check applies only when a local GGUF window is known; + // external providers have ggufContextLength === null. + const withinLocalLimit = + !store.ggufContextLength || + (savedUsage?.totalTokens ?? 0) <= store.ggufContextLength; + // Legacy unscoped usage (no modelId) is only trusted when a + // known local window bounds the totals, so we can't misattribute + // an old local turn to a newly-selected external provider. + const modelMatches = savedUsage?.modelId + ? savedUsage.modelId === store.params.checkpoint + : typeof store.ggufContextLength === "number" && + store.ggufContextLength > 0; + if (savedUsage && withinLocalLimit && modelMatches) { store.setContextUsage(savedUsage); } 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 a00b53a44c..6b60ed51ea 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -291,6 +291,8 @@ type ChatRuntimeStore = { completionTokens: number; totalTokens: number; cachedTokens: number; + // Anthropic-only; optional so pre-cache-stats persisted entries load. + cacheWriteTokens?: number; } | null; modelLoading: boolean; activeNativePathToken: string | null; @@ -640,7 +642,14 @@ export const useChatRuntimeStore = create((set, get) => ({ if (state.settingsHydrated && hasKeys(changedParams)) { saveSettingsPatch({ inferenceParams: changedParams }); } - return { params }; + // Mirror setCheckpoint: the local model load path can mutate + // params.checkpoint via setParams() before setCheckpoint runs, + // leaving stale per-turn counters under the new checkpoint. + const checkpointChanged = state.params.checkpoint !== params.checkpoint; + return { + params, + ...(checkpointChanged ? { contextUsage: null } : {}), + }; }), setCustomPresets: (customPresets) => set(() => { @@ -704,12 +713,17 @@ export const useChatRuntimeStore = create((set, get) => ({ // mount, and a stale persisted local id would race against the // freshly-loaded model. See LAST_EXTERNAL_CHECKPOINT_KEY notes. saveLastExternalCheckpoint(isExternalModelId(modelId) ? modelId : null); + // Clear stale per-turn usage when the model changes; the relaxed + // external-provider render gate would otherwise show old counters + // until the next completion overwrites them. + const checkpointChanged = state.params.checkpoint !== modelId; return { params: { ...state.params, checkpoint: modelId, }, activeGgufVariant: ggufVariant ?? null, + ...(checkpointChanged ? { contextUsage: null } : {}), }; }), setActiveThreadId: (activeThreadId) => From 7d1b68079ed8d9fec9a800346a2f450456784132 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 23:37:12 -0700 Subject: [PATCH 11/43] Studio: Anthropic fast_mode toggle and streaming refusal handling (#5715) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: add Anthropic fast_mode toggle + surface streaming refusals Fast mode (beta `fast-mode-2026-02-01`) lets Claude Opus 4.6 and 4.7 generate output tokens up to 2.5x faster at 6x standard Opus pricing. The toggle lives in Configuration → Provider when the selected Anthropic model is Opus 4.6 or 4.7 and is otherwise hidden. Backend gates the same prefixes a second time so a stale frontend cannot make Anthropic 400 the request, and the `fast-mode-2026-02-01` beta header is merged onto whatever other betas the request already needed (code-execution, compaction). Streaming refusals (`message_delta.delta.stop_reason="refusal"` on Claude 4 models) now surface a short user-facing notice in the assistant message before the translated OpenAI chunk emits the existing `finish_reason="content_filter"`. Previously the chat bubble truncated silently because the SSE stopped mid-stream with no visible explanation. Per the upstream docs the conversation must be reset before continuing, so the notice tells the user exactly that. Reference: - https://platform.claude.com/docs/en/build-with-claude/fast-mode - https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals Tests: - studio/backend/tests/test_anthropic_fast_mode_and_refusal.py (8 cases pinning fast_mode pass-through on 4.6/4.7, silent drop on Sonnet / Haiku / older Opus / None / False, and the refusal notice + finish reason on a synthetic refusal stream). * Studio: drop refused Anthropic turns from the next request Anthropic's streaming-refusal guidance says the refused assistant turn must be removed or updated before the next call -- otherwise the safety classifier keeps refusing. The PR only added a user-visible notice; the partial assistant output (plus the notice itself) still rode the next request via toOpenAIMessage. Tag the refusal turn with an HTML-comment sentinel emitted alongside the notice. The chat-adapter checks for that sentinel in toOpenAIMessage and returns null, so the refused turn is excluded from outboundMessages. The notice still renders in the transcript (HTML comments don't display), so users keep the explanation. * Studio: filter None finish_reason entries in test helper test_refusal_maps_to_content_filter expects only ['content_filter'] in the finish_reasons list, but the post-PR refusal path emits a user-visible content notice chunk first. Every _content_chunk carries 'finish_reason: None' by construction; the helper was appending those, so the assertion saw [None, 'content_filter'] instead of ['content_filter']. None is not a finish reason -- it's just mid-stream delta noise. Skip None values in _finish_reasons so the helper reflects what the test names actually claim to check. Same fix applies cleanly to the other helper usages (pause_turn test expects [] and the sibling stop test expects ['stop'], both unaffected). * Studio: cover Anthropic fast-mode edge cases Adds 19 cases on top of the 9 in test_anthropic_fast_mode_and_refusal. The base file pins the happy path; this file fills in the cliffs: * Dated-snapshot prefix matching: claude-opus-4-7-2026-02-01 and claude-opus-4-6-2026-02-01 still gate fast_mode through, while claude-opus-4-5-2025-08-01 and claude-sonnet-4-6-2026-02-01 do not. * Strict opt-in: a future claude-opus-4-8 or claude-opus-5 does NOT auto-enable fast_mode -- the prefix tuple must be bumped explicitly when a new family is whitelisted upstream. * Beta-header merge: fast_mode coexists with code-execution-2025-08-25 and compact-2026-01-12 in one comma-separated anthropic-beta header with no duplicates and no truncation. Pins the value to the exact fast-mode-2026-02-01 docs token so a typo would fail CI. * Non-destruction: fast_mode=None produces byte-identical outbound body and headers to the version that omits the argument entirely. Same for fast_mode=False. Guarantees the upgrade path is non-breaking on existing Anthropic streams. * Refusal stream ordering: the user-visible notice precedes the finish_reason chunk so a streaming UI paints text before flipping to content_filter. Refusal sentinel emitted exactly once. Notice rides a normal content delta chunk with finish_reason still null. Partial assistant deltas survive before the notice. * Provider-side refusal coverage: a refusal on Sonnet (not just Opus) still emits the notice + sentinel + content_filter mapping, since refusal handling is not gated on fast-mode capability. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Persist fastMode, drop refused user message on retry Two follow-ups on #5715: 1) sanitizeInferenceParams stripped fastMode. fastMode is in PERSISTED_INFERENCE_PARAM_KEYS but the storage sanitizer only kept numeric fields plus systemPrompt and trustRemoteCode, so the new toggle was silently dropped on reload and on the /api/chat/settings round-trip. Save it the same way trustRemoteCode is saved. 2) Refusal recovery now also drops the triggering user turn. Returning null from toOpenAIMessage on the assistant side left the user prompt that caused the refusal in the outbound history, so the very next request would re-trigger the same classifier. Anthropic's refusal-handling guidance is explicit on this: remove the refused turn AND the user message that triggered it before the next call. Implemented via a pre-pass that pops the trailing user message when an assistant carries the refusal sentinel. Typecheck clean. * Studio: out-of-band refusal signal + fast-mode prefix/usage/pricing fixes The text sentinel for the Anthropic refusal drop signal was spoofable: any assistant message containing the literal would prune the prior user + assistant pair on the next request. Move the signal onto a separate _toolEvent chunk that the chat adapter latches into assistant.metadata.custom.anthropicRefusal; assistant text can no longer control the pruner. Tighten the fast-mode model gate (backend + frontend) to require a "-" family boundary so claude-opus-4-70 / claude-opus-4-7b style IDs do not get speed: "fast" on a naive startswith match. Use survivingMessages for the image / audio attachment scan so a refused user turn does not gate or mis-attribute the next non-refused turn. Propagate Anthropic usage.speed onto the OpenAI-style usage chunk and apply the documented 6x fast-mode multiplier in the cost calculator (stacks with prompt-cache multipliers per the docs); expose the new multiplier on the pricing snapshot for the UI tooltip. Tests cover the tool-event chunk shape, the prefix-collision rejects, usage.speed propagation, the 6x pricing math, and that the visible refusal text carries no embedded sentinel. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shorten fast-mode and refusal comments for PR #5715 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 63 +++ studio/backend/core/inference/pricing.py | 13 + studio/backend/models/inference.py | 10 + studio/backend/routes/inference.py | 1 + .../test_anthropic_fast_mode_and_refusal.py | 164 +++++++ .../tests/test_anthropic_fast_mode_edge.py | 442 ++++++++++++++++++ .../backend/tests/test_anthropic_web_fetch.py | 9 +- studio/backend/tests/test_pricing.py | 60 +++ .../src/features/chat/api/chat-adapter.ts | 67 ++- .../src/features/chat/chat-settings-sheet.tsx | 29 ++ .../features/chat/provider-capabilities.ts | 24 + .../chat/stores/chat-runtime-store.ts | 1 + .../frontend/src/features/chat/types/api.ts | 6 + .../src/features/chat/types/runtime.ts | 7 + .../chat/utils/chat-settings-storage.ts | 5 + 15 files changed, 894 insertions(+), 7 deletions(-) create mode 100644 studio/backend/tests/test_anthropic_fast_mode_and_refusal.py create mode 100644 studio/backend/tests/test_anthropic_fast_mode_edge.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 25e1725337..d8a36610b8 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -173,10 +173,28 @@ _ANTHROPIC_COMPACTION_TYPE = "compact_20260112" _ANTHROPIC_COMPACTION_MIN = 50_000 +# Anthropic fast-mode beta (Opus 4.6 / 4.7 only, per +# https://platform.claude.com/docs/en/build-with-claude/fast-mode). +# Mutually exclusive with the Priority service tier. +_ANTHROPIC_FAST_MODE_BETA = "fast-mode-2026-02-01" +_ANTHROPIC_FAST_MODE_PREFIXES = ( + "claude-opus-4-7", + "claude-opus-4-6", +) + + def _anthropic_supports_compaction(model: str) -> bool: return model.startswith(_ANTHROPIC_COMPACTION_PREFIXES) +def _anthropic_supports_fast_mode(model: str) -> bool: + # Require a family boundary ("" or "-") after the prefix so IDs like + # "claude-opus-4-70" / "claude-opus-4-7b" do not match. + return any( + model == p or model.startswith(f"{p}-") for p in _ANTHROPIC_FAST_MODE_PREFIXES + ) + + class _MistralThinkingSpec(NamedTuple): models: tuple[str, ...] style: Literal["prompt_mode", "reasoning_effort", "disabled"] @@ -348,6 +366,7 @@ class ExternalProviderClient: anthropic_code_exec_container_id: Optional[str] = None, prompt_cache_ttl: Optional[str] = None, compaction_threshold: Optional[int] = None, + fast_mode: Optional[bool] = None, stream: bool = True, ) -> AsyncGenerator[str, None]: """ @@ -360,6 +379,9 @@ class ExternalProviderClient: supplies a value the provider accepts — the frontend's provider-capability map already filters these per provider, so we treat them as opt-in here. + + ``fast_mode`` only applies to Anthropic Opus 4.6 / 4.7 (silently + dropped elsewhere); adds the beta header and ``speed: "fast"``. """ if not self._is_openai_compatible(): async for line in self._stream_anthropic( @@ -376,6 +398,7 @@ class ExternalProviderClient: anthropic_code_exec_container_id, prompt_cache_ttl, compaction_threshold, + fast_mode = fast_mode, ): yield line return @@ -1186,6 +1209,8 @@ class ExternalProviderClient: anthropic_code_exec_container_id: Optional[str] = None, prompt_cache_ttl: Optional[str] = None, compaction_threshold: Optional[int] = None, + *, + fast_mode: Optional[bool] = None, ) -> AsyncGenerator[str, None]: """ Call the Anthropic Messages API and translate its SSE to OpenAI format. @@ -1611,6 +1636,13 @@ class ExternalProviderClient: ] } + # fast_mode is Opus 4.6/4.7 only; silently drop elsewhere. + # Incompatible with the Priority service_tier (frontend gate + # prevents both at once; backend lets Anthropic 400 if combined). + fast_mode_active = bool(fast_mode) and _anthropic_supports_fast_mode(model) + if fast_mode_active: + body["speed"] = "fast" + url = f"{self.base_url}/messages" completion_id = f"chatcmpl-anthropic-{model.replace('/', '-')}" @@ -1669,6 +1701,8 @@ class ExternalProviderClient: beta_parts.append(_ANTHROPIC_CODE_EXECUTION_BETA) if compaction_active and _ANTHROPIC_COMPACTION_BETA not in beta_parts: beta_parts.append(_ANTHROPIC_COMPACTION_BETA) + if fast_mode_active and _ANTHROPIC_FAST_MODE_BETA not in beta_parts: + beta_parts.append(_ANTHROPIC_FAST_MODE_BETA) if beta_parts: request_headers["anthropic-beta"] = ",".join(beta_parts) @@ -2410,6 +2444,29 @@ class ExternalProviderClient: # finish_reason="stop" chunk that would # truncate the rendered message in the UI. mapped = _finish_reason_map.get(stop_reason, "stop") + # Streaming refusal: emit a visible notice + # plus an out-of-band _toolEvent so the + # frontend can prune the refused turn. + # The mapped finish_reason is + # "content_filter" per OpenAI spec. + # https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals + if stop_reason == "refusal": + logger.warning( + "Anthropic refusal stop_reason (model=%s)", + model, + ) + # Drop signal rides _toolEvent (not + # text) so assistant content cannot + # spoof a context reset. + yield _content_chunk( + "\n\n_The response was stopped by " + "Anthropic's safety classifier. Edit " + "or remove the previous turn and try " + "again._" + ) + yield _emit_tool_event( + {"type": "anthropic_refusal"} + ) if mapped is not None: chunk = { "id": completion_id, @@ -3939,6 +3996,12 @@ def _build_usage_chunk( "cache_creation_input_tokens": cache_creation, "cache_read_input_tokens": cache_read, } + # Propagate fast-mode `usage.speed` so the cost ledger can apply + # the 6x multiplier without re-derivation (Anthropic falls back + # to "standard" when fast-mode is unsupported or rate-limited). + speed = last_usage.get("speed") + if speed in ("fast", "standard"): + usage_block["speed"] = speed else: prompt_tokens = last_usage.get("input_tokens") or 0 cached = 0 diff --git a/studio/backend/core/inference/pricing.py b/studio/backend/core/inference/pricing.py index 74c57fa594..4241f67296 100644 --- a/studio/backend/core/inference/pricing.py +++ b/studio/backend/core/inference/pricing.py @@ -105,6 +105,9 @@ OPENAI_PRICING: dict[str, dict[str, float]] = { ANTHROPIC_CACHE_5M_WRITE_MULT = 1.25 ANTHROPIC_CACHE_1H_WRITE_MULT = 2.0 ANTHROPIC_CACHE_READ_MULT = 0.1 +# Anthropic fast-mode (Opus 4.6 / 4.7 only): 6x standard on input + output. +# https://platform.claude.com/docs/en/build-with-claude/fast-mode#pricing +ANTHROPIC_FAST_MODE_MULT = 6.0 # OpenAI: cache reads are 0.1x base input, cache writes are not billed # separately (the first prefix-write request just pays normal input). @@ -235,6 +238,15 @@ def calculate_cost( base = prices["input_per_mtok"] out_per = prices["output_per_mtok"] + # Anthropic fast-mode: 6x on input + output. Cache multipliers stack + # on top of fast-mode, so applying once to (base, out_per) propagates + # into the cache_*_usd buckets computed below. + if provider == "anthropic" and usage.get("speed") == "fast": + base *= ANTHROPIC_FAST_MODE_MULT + out_per *= ANTHROPIC_FAST_MODE_MULT + if out["model_priced"]: + out["model_priced"] = f"{out['model_priced']} (fast)" + out["input_usd"] = (input_tokens / 1_000_000.0) * base out["output_usd"] = (output_tokens / 1_000_000.0) * out_per @@ -315,6 +327,7 @@ def pricing_snapshot() -> dict[str, Any]: "cache_5m_write_mult": ANTHROPIC_CACHE_5M_WRITE_MULT, "cache_1h_write_mult": ANTHROPIC_CACHE_1H_WRITE_MULT, "cache_read_mult": ANTHROPIC_CACHE_READ_MULT, + "fast_mode_mult": ANTHROPIC_FAST_MODE_MULT, "web_search_usd_per_1k": ANTHROPIC_WEB_SEARCH_USD_PER_1K, "code_execution_usd_per_hour": ANTHROPIC_CODE_EXEC_USD_PER_HOUR, }, diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b5626951c4..68bc7a7017 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -786,6 +786,16 @@ class ChatCompletionRequest(BaseModel): "to auto-create." ), ) + fast_mode: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Anthropic fast-mode toggle. On Claude Opus 4.6 / " + "4.7 adds the `fast-mode-2026-02-01` beta header and sends " + "`speed: 'fast'` for higher OTPS at premium pricing. Silently " + "ignored on every other model + provider. See " + "https://platform.claude.com/docs/en/build-with-claude/fast-mode" + ), + ) @model_validator(mode = "after") def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest": diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index bf92055929..143947efc8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1876,6 +1876,7 @@ async def _proxy_to_external_provider( anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id, prompt_cache_ttl = payload.prompt_cache_ttl, compaction_threshold = payload.compaction_threshold, + fast_mode = payload.fast_mode, stream = payload.stream, ) try: diff --git a/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py b/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py new file mode 100644 index 0000000000..e7e5ec64d4 --- /dev/null +++ b/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py @@ -0,0 +1,164 @@ +# 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 Anthropic fast-mode wiring and streaming refusal handling. + +fast_mode=True on Opus 4.6/4.7 attaches the ``fast-mode-2026-02-01`` +beta header and sets ``speed: "fast"``; unsupported models drop both. +Streaming ``stop_reason: "refusal"`` surfaces a user notice before the +``content_filter`` finish chunk. +https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _empty_message_sse() -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"claude-opus-4-7","role":"assistant",' + b'"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"end_turn"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _refusal_sse() -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"claude-opus-4-7","role":"assistant",' + b'"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"Hello."}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"refusal"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]: + """Install a MockTransport, drive one streamed call, return body+lines.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = sse or _empty_message_sse(), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + out_lines: list[str] = [] + + async def run(): + client = _make_client() + try: + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = kwargs.get("model", "claude-opus-4-7"), + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + fast_mode = kwargs.get("fast_mode"), + ): + out_lines.append(line) + finally: + await client.close() + + _drive(run()) + return captured, out_lines + + +def test_fast_mode_attaches_beta_header_and_speed_on_opus_4_7(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7") + assert cap["body"].get("speed") == "fast", cap["body"] + beta = cap["headers"].get("anthropic-beta", "") + assert "fast-mode-2026-02-01" in beta, beta + + +def test_fast_mode_attaches_beta_header_and_speed_on_opus_4_6(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-6") + assert cap["body"].get("speed") == "fast", cap["body"] + assert "fast-mode-2026-02-01" in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_dropped_on_sonnet(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-sonnet-4-6") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_dropped_on_haiku(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-haiku-4-5") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_dropped_on_older_opus(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-5") + assert "speed" not in cap["body"], cap["body"] + + +def test_fast_mode_false_does_not_attach_header_or_field(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = False) + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_none_does_not_attach_header_or_field(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = None) + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_refusal_emits_user_facing_notice_and_content_filter_finish(monkeypatch): + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + # User-visible refusal notice. + assert "stopped by Anthropic's safety classifier" in body, body + # OpenAI-spec finish_reason mapping. + assert '"finish_reason": "content_filter"' in body, body + # Original deltas preserved before the refusal supplement. + assert "Hello." in body, body + + +def test_refusal_emits_tool_event_for_chat_adapter_drop(monkeypatch): + """Refused turns emit an out-of-band `_toolEvent` that the chat-adapter + latches into assistant `metadata.custom.anthropicRefusal`, driving + the next-request prune. Tool event (not text) prevents spoofing. + """ + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + assert '"_toolEvent": {"type": "anthropic_refusal"}' in body, body + # Visible refusal text must not embed a sentinel that could spoof + # a context reset if echoed by another assistant message. + assert "studio:anthropic-refusal" not in body, body diff --git a/studio/backend/tests/test_anthropic_fast_mode_edge.py b/studio/backend/tests/test_anthropic_fast_mode_edge.py new file mode 100644 index 0000000000..0052cb94ad --- /dev/null +++ b/studio/backend/tests/test_anthropic_fast_mode_edge.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case coverage for the Anthropic fast-mode + refusal wiring. + +Complements ``test_anthropic_fast_mode_and_refusal.py`` (happy path) +with dated snapshots, strict opt-in (future Opus families do not +auto-enable), multi-beta header merging, refusal stream ordering, and +the non-destruction guarantee for unset/None fast_mode. +""" + +import asyncio +import json +import re + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _empty_message_sse(model: str = "claude-opus-4-7") -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"' + model.encode() + b'",' + b'"role":"assistant","stop_reason":null,"usage":' + b'{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"end_turn"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _refusal_sse(model: str = "claude-opus-4-7") -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"' + model.encode() + b'",' + b'"role":"assistant","stop_reason":null,"usage":' + b'{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"Hello."}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"refusal"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]: + """Install a MockTransport, drive one streamed call, return body+lines.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = sse or _empty_message_sse(kwargs.get("model", "claude-opus-4-7")), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + out_lines: list[str] = [] + + async def run(): + client = _make_client() + try: + extra = {} + for key in ( + "enabled_tools", + "compaction_threshold", + "fast_mode", + ): + if key in kwargs: + extra[key] = kwargs[key] + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = kwargs.get("model", "claude-opus-4-7"), + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + **extra, + ): + out_lines.append(line) + finally: + await client.close() + + _drive(run()) + return captured, out_lines + + +# ──────────────────────────── dated snapshot prefix ──────────────────────────── +def test_fast_mode_attaches_on_dated_opus_4_7_snapshot(monkeypatch): + """Dated snapshot ``claude-opus-4-7-2026-02-01`` must match the prefix.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7-2026-02-01") + assert cap["body"].get("speed") == "fast", cap["body"] + assert "fast-mode-2026-02-01" in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_attaches_on_dated_opus_4_6_snapshot(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-6-2026-02-01") + assert cap["body"].get("speed") == "fast", cap["body"] + assert "fast-mode-2026-02-01" in cap["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── strict opt-in semantics ──────────────────────────── +def test_fast_mode_does_not_auto_enable_on_future_opus_4_8(monkeypatch): + """Future ``claude-opus-4-8`` must not auto-enable; opt-in per family.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-8") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_does_not_auto_enable_on_future_opus_5(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-5") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_does_not_auto_enable_on_sonnet_dated_snapshot(monkeypatch): + """Sonnet snapshots share the compaction prefix but not fast_mode.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-sonnet-4-6-2026-02-01") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── beta header merge ──────────────────────────── +def _beta_parts(headers: dict) -> list[str]: + raw = headers.get("anthropic-beta", "") + return [p.strip() for p in raw.split(",") if p.strip()] + + +def test_fast_mode_merges_with_code_execution_beta(monkeypatch): + """fast_mode + code_execution -> two comma-separated betas, no overwrite.""" + cap, _ = _capture( + monkeypatch, + fast_mode = True, + model = "claude-opus-4-7", + enabled_tools = ["code_execution"], + ) + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts, cap["headers"] + assert any(p.startswith("code-execution-") for p in parts), cap["headers"] + # No duplicates. + assert len(parts) == len(set(parts)), parts + + +def test_fast_mode_merges_with_compaction_beta(monkeypatch): + """fast_mode + compaction_threshold >= 50K -> both betas present.""" + cap, _ = _capture( + monkeypatch, + fast_mode = True, + model = "claude-opus-4-7", + compaction_threshold = 100_000, + ) + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts, cap["headers"] + assert "compact-2026-01-12" in parts, cap["headers"] + + +def test_fast_mode_merges_with_code_execution_and_compaction(monkeypatch): + """Three betas coexist in one comma-separated header, no duplicates.""" + cap, _ = _capture( + monkeypatch, + fast_mode = True, + model = "claude-opus-4-7", + enabled_tools = ["code_execution"], + compaction_threshold = 100_000, + ) + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts + assert "compact-2026-01-12" in parts + assert any(p.startswith("code-execution-") for p in parts), parts + assert len(parts) >= 3 + assert len(parts) == len(set(parts)), parts + + +def test_fast_mode_beta_value_is_pinned(monkeypatch): + """Pin the exact beta tag ``fast-mode-2026-02-01`` from the docs.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7") + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts, parts + # Reject obvious typos. + assert not any(p.startswith("fastmode-") for p in parts), parts + assert not any("fast_mode" in p for p in parts), parts + + +# ──────────────────────────── non-destruction guarantee ──────────────────────────── +def test_fast_mode_unset_is_byte_identical_to_omitted(monkeypatch): + """``fast_mode=None`` must produce the same body/headers as omission.""" + cap_none, _ = _capture(monkeypatch, fast_mode = None, model = "claude-opus-4-7") + + # Re-run without passing fast_mode at all. + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = _empty_message_sse(), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + async def run(): + client = _make_client() + try: + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + ): + pass + finally: + await client.close() + + _drive(run()) + + assert cap_none["body"] == captured["body"], (cap_none["body"], captured["body"]) + # Headers can vary by httpx-injected fields (host, connection); compare + # the load-bearing ones. + for key in ("anthropic-version", "x-api-key", "content-type"): + assert cap_none["headers"].get(key) == captured["headers"].get(key), key + assert "anthropic-beta" not in cap_none["headers"] + assert "anthropic-beta" not in captured["headers"] + assert "speed" not in cap_none["body"] + assert "speed" not in captured["body"] + + +def test_fast_mode_false_on_opus_4_7_byte_identical_to_unset(monkeypatch): + """``fast_mode=False`` produces the same outbound shape as unset.""" + cap_false, _ = _capture(monkeypatch, fast_mode = False, model = "claude-opus-4-7") + assert "speed" not in cap_false["body"], cap_false["body"] + assert "fast-mode-2026-02-01" not in cap_false["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── refusal stream ordering ──────────────────────────── +def test_refusal_notice_appears_before_content_filter_chunk(monkeypatch): + """The notice content delta must precede the finish_reason chunk.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + notice_idx = next(i for i, l in enumerate(lines) if "stopped by Anthropic" in l) + filter_idx = next( + i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l + ) + assert notice_idx < filter_idx, (notice_idx, filter_idx, lines) + + +def test_refusal_tool_event_emitted_exactly_once(monkeypatch): + """A single refusal emits the chat-adapter drop signal exactly once.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + count = body.count('"_toolEvent": {"type": "anthropic_refusal"}') + assert count == 1, (count, body) + + +def test_refusal_text_carries_no_html_sentinel(monkeypatch): + """Visible refusal text must not embed a ``studio:anthropic-refusal`` + sentinel; the drop signal rides _toolEvent only.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + assert "studio:anthropic-refusal" not in body, body + + +def test_refusal_handling_works_on_sonnet_model(monkeypatch): + """Refusal handling is provider-side; Sonnet refusals must also surface.""" + _, lines = _capture( + monkeypatch, sse = _refusal_sse("claude-sonnet-4-6"), model = "claude-sonnet-4-6" + ) + body = "\n".join(lines) + assert "stopped by Anthropic's safety classifier" in body, body + assert '"_toolEvent": {"type": "anthropic_refusal"}' in body, body + assert '"finish_reason": "content_filter"' in body, body + + +def test_refusal_preserves_partial_assistant_text(monkeypatch): + """Partial deltas already streamed must precede the refusal notice.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + body = "\n".join(lines) + hello_idx = body.index("Hello.") + notice_idx = body.index("stopped by Anthropic") + assert hello_idx < notice_idx, (hello_idx, notice_idx) + + +def test_refusal_chunk_is_proper_openai_delta_shape(monkeypatch): + """The notice rides ``choices[0].delta.content`` (not a finish chunk); + OpenAI-spec clients treat it as ordinary streamed text.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + # Find the chunk that carries the refusal text. + notice_chunk = None + for line in lines: + if line.startswith("data: ") and "stopped by Anthropic" in line: + notice_chunk = json.loads(line[len("data: ") :]) + break + assert notice_chunk is not None, lines + choice = notice_chunk["choices"][0] + assert "delta" in choice and "content" in choice["delta"], notice_chunk + # Must NOT carry a finish_reason itself -- that comes on the next + # chunk. + assert choice.get("finish_reason") in (None,), notice_chunk + # Refusal text is plain-spoken; no embedded sentinel. + assert "studio:anthropic-refusal" not in choice["delta"]["content"] + + +def test_refusal_tool_event_chunk_shape(monkeypatch): + """Drop signal rides a Studio `_toolEvent` envelope (delta={}, + finish_reason=null); the frontend latches on + `_toolEvent.type == "anthropic_refusal"`.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + refusal_chunk = None + for line in lines: + if line.startswith("data: ") and "anthropic_refusal" in line: + refusal_chunk = json.loads(line[len("data: ") :]) + break + assert refusal_chunk is not None, lines + assert refusal_chunk["_toolEvent"] == {"type": "anthropic_refusal"}, refusal_chunk + choice = refusal_chunk["choices"][0] + assert choice["delta"] == {}, refusal_chunk + assert choice["finish_reason"] is None, refusal_chunk + + +# ──────────────────────────── future-proofing ──────────────────────────── +def test_fast_mode_prefix_tuple_matches_capability_doc(monkeypatch): + """Tuple must exactly match the two families in the upstream docs: + https://platform.claude.com/docs/en/build-with-claude/fast-mode.""" + from core.inference.external_provider import _ANTHROPIC_FAST_MODE_PREFIXES + + assert set(_ANTHROPIC_FAST_MODE_PREFIXES) == { + "claude-opus-4-7", + "claude-opus-4-6", + }, _ANTHROPIC_FAST_MODE_PREFIXES + + +def test_fast_mode_speed_field_value_is_literal_fast(monkeypatch): + """Pin the wire value to the literal string ``"fast"``.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7") + assert cap["body"]["speed"] == "fast", cap["body"] + + +def test_fast_mode_dropped_on_opus_4_5_dated_snapshot(monkeypatch): + """Previous-family snapshots like ``claude-opus-4-5-2025-...`` must not match.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-5-2025-08-01") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_rejects_prefix_collision_4_70(monkeypatch): + """IDs like ``claude-opus-4-70`` / ``-4-7b`` must not match the prefix.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-70") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_rejects_prefix_collision_4_7b(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7b") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_rejects_prefix_collision_4_6_extra(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-60") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── usage.speed propagation ──────────────────────────── +def _fast_speed_sse(model: str = "claude-opus-4-7", speed: str = "fast") -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"' + model.encode() + b'",' + b'"role":"assistant","stop_reason":null,"usage":' + b'{"input_tokens":4,"output_tokens":1}}}\n\n' + b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"hi"}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"end_turn"},' + b'"usage":{"output_tokens":5,"speed":"' + speed.encode() + b'"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch): + """``usage.speed == "fast"`` from upstream must reach the Studio usage chunk.""" + _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "fast")) + usage_lines = [l for l in lines if l.startswith("data: ") and '"usage"' in l] + assert usage_lines, lines + parsed = [json.loads(l[len("data: ") :]) for l in usage_lines] + speeds = [p["usage"].get("speed") for p in parsed if "usage" in p] + assert "fast" in speeds, parsed + + +def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch): + _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "standard")) + parsed = [ + json.loads(l[len("data: ") :]) + for l in lines + if l.startswith("data: ") and '"usage"' in l + ] + speeds = [p["usage"].get("speed") for p in parsed if "usage" in p] + assert "standard" in speeds, parsed + + +def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch): + """Studio must not invent ``usage.speed`` when upstream omits it.""" + _, lines = _capture(monkeypatch) + parsed = [ + json.loads(l[len("data: ") :]) + for l in lines + if l.startswith("data: ") and '"usage"' in l + ] + for p in parsed: + usage = p.get("usage") or {} + assert "speed" not in usage, p diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index cdb5f6254c..7277757fcf 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -365,7 +365,9 @@ def test_web_fetch_error_renders_error_code(monkeypatch): def _finish_reasons(lines: list[str]) -> list: - """Return the finish_reason fields from every chat.completion.chunk.""" + """Return non-null finish_reason fields from each chat.completion.chunk. + Mid-stream content deltas carry ``finish_reason: None`` and are skipped + (the refusal path emits a notice delta before the content_filter chunk).""" out: list = [] for line in lines: if not line.startswith("data:"): @@ -380,8 +382,9 @@ def _finish_reasons(lines: list[str]) -> list: if parsed.get("object") != "chat.completion.chunk": continue for choice in parsed.get("choices") or []: - if "finish_reason" in choice: - out.append(choice["finish_reason"]) + reason = choice.get("finish_reason") + if reason is not None: + out.append(reason) return out diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py index cc8c16993c..f534f1f58a 100644 --- a/studio/backend/tests/test_pricing.py +++ b/studio/backend/tests/test_pricing.py @@ -14,6 +14,7 @@ from core.inference.pricing import ( ANTHROPIC_CACHE_5M_WRITE_MULT, ANTHROPIC_CACHE_1H_WRITE_MULT, ANTHROPIC_CACHE_READ_MULT, + ANTHROPIC_FAST_MODE_MULT, ANTHROPIC_PRICING, OPENAI_CACHE_READ_MULT, OPENAI_CONTAINER_USD_PER_HOUR, @@ -57,6 +58,64 @@ def test_anthropic_opus_4_7_input_and_output_math(): assert _isclose(out["total_usd"], 30.0) +# ── Anthropic fast-mode 6x multiplier (Opus 4.6 / 4.7 only) ───────── + + +def test_anthropic_fast_mode_charges_6x_standard_opus(): + """6x on input + output when ``usage.speed == "fast"``. + https://platform.claude.com/docs/en/build-with-claude/fast-mode""" + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 1_000_000, + "output_tokens": 1_000_000, + "speed": "fast", + }, + ) + assert _isclose(out["input_usd"], 5.0 * ANTHROPIC_FAST_MODE_MULT) + assert _isclose(out["output_usd"], 25.0 * ANTHROPIC_FAST_MODE_MULT) + assert _isclose(out["total_usd"], 30.0 * ANTHROPIC_FAST_MODE_MULT) + assert "(fast)" in out["model_priced"], out["model_priced"] + + +def test_anthropic_fast_mode_does_not_affect_standard_speed(): + """``speed: "standard"`` (or missing) keeps the base rates.""" + out_standard = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 1_000_000, + "output_tokens": 1_000_000, + "speed": "standard", + }, + ) + out_missing = calculate_cost( + "anthropic", + "claude-opus-4-7", + {"input_tokens": 1_000_000, "output_tokens": 1_000_000}, + ) + assert _isclose(out_standard["total_usd"], out_missing["total_usd"]) + assert _isclose(out_standard["total_usd"], 30.0) + + +def test_anthropic_fast_mode_stacks_with_cache_read_multiplier(): + """Cache multipliers apply on top of fast-mode (per docs).""" + base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 0, + "output_tokens": 0, + "cache_read_input_tokens": 1_000_000, + "speed": "fast", + }, + ) + expected = base * ANTHROPIC_FAST_MODE_MULT * ANTHROPIC_CACHE_READ_MULT + assert _isclose(out["cache_read_usd"], expected) + + # ── Anthropic cache write 5m + read multipliers ────────────────────── @@ -412,6 +471,7 @@ def test_snapshot_contains_provider_buckets_and_multipliers(): assert a["cache_5m_write_mult"] == ANTHROPIC_CACHE_5M_WRITE_MULT assert a["cache_1h_write_mult"] == ANTHROPIC_CACHE_1H_WRITE_MULT assert a["cache_read_mult"] == ANTHROPIC_CACHE_READ_MULT + assert a["fast_mode_mult"] == ANTHROPIC_FAST_MODE_MULT assert "web_search_usd_per_1k" in a assert "code_execution_usd_per_hour" in a assert "models" in o and "gpt-5.5" in o["models"] diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 9842e380e0..800395bddc 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -28,6 +28,7 @@ import { providerSupportsBuiltinImageGeneration, providerSupportsBuiltinWebFetch, providerSupportsBuiltinWebSearch, + providerSupportsFastMode, } from "../provider-capabilities"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; import { useExternalProvidersStore } from "../stores/external-providers-store"; @@ -347,6 +348,18 @@ function collectImageParts( return parts; } +// Refusal flag stamped on assistant metadata when the backend emits the +// `anthropic_refusal` _toolEvent. We drop the refused pair from the next +// request body (Anthropic guidance: leaving refusals in context keeps +// refusing). Metadata (not text) prevents content from spoofing a reset. +function isAnthropicRefusalMessage(message: RunMessage): boolean { + if (message.role !== "assistant") return false; + const metadata = (message as { metadata?: unknown }).metadata as + | { custom?: Record } + | undefined; + return metadata?.custom?.anthropicRefusal === true; +} + function toOpenAIMessage(message: RunMessage): { role: "system" | "user" | "assistant"; content: OpenAIMessageContent; @@ -367,6 +380,11 @@ function toOpenAIMessage(message: RunMessage): { /data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g, "[audio]", ); + if (isAnthropicRefusalMessage(message)) { + // Prune refused assistant turn from outbound history; the + // rendered transcript still shows the user-visible notice. + return null; + } } const imageParts = collectImageParts(message); @@ -925,7 +943,24 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ), ); - const outboundMessages = messages + // Two-pass build: a refused assistant turn also drops the user + // prompt that triggered it (leaving it in context re-triggers + // the classifier). Refusal flag rides assistant + // metadata.custom.anthropicRefusal, set out-of-band from the + // backend _toolEvent. + const survivingMessages: RunMessage[] = []; + for (const message of messages) { + if (isAnthropicRefusalMessage(message)) { + const last = survivingMessages.at(-1); + if (last && last.role === "user") { + survivingMessages.pop(); + } + continue; + } + survivingMessages.push(message); + } + + const outboundMessages = survivingMessages .map(toOpenAIMessage) .filter((message): message is NonNullable => Boolean(message), @@ -995,8 +1030,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }); } } - const imageBase64 = findLatestUserImageBase64(messages); - const audioBase64 = findLatestUserAudioBase64(messages); + // Scan post-prune history so a refused user turn's image/audio + // doesn't gate or mis-attribute the next non-refused turn. + const imageBase64 = findLatestUserImageBase64(survivingMessages); + const audioBase64 = findLatestUserAudioBase64(survivingMessages); // Block when ANY image is in the outbound payload (current or // prior turns) and the loaded model can't process images. Keeps @@ -1032,7 +1069,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (audioBase64) { const audioName = runtime.pendingAudioName; if (audioName) { - const lastUserMsg = [...messages] + const lastUserMsg = [...survivingMessages] .reverse() .find((m) => m.role === "user"); if (lastUserMsg) sentAudioNames.set(lastUserMsg.id, audioName); @@ -1143,6 +1180,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Tool call content parts — accumulated and yielded cumulatively. // result is set directly on the tool-call part when tool_end arrives. const toolCallParts: ToolCallMessagePart[] = []; + // Latched on the `anthropic_refusal` tool event; stamped onto the + // final assistant metadata as `custom.anthropicRefusal` to drive + // the history-prune above. + let anthropicRefusalSeen = false; let serverMetadata: { usage?: ServerUsage; timings?: ServerTimings; @@ -1485,6 +1526,16 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { isPromptCacheTtl(externalProvider.promptCacheTtl) ? { prompt_cache_ttl: externalProvider.promptCacheTtl } : {}), + // Anthropic fast mode (Opus 4.6 / 4.7 only); backend + // silently drops on unsupported models as a second + // line of defence. + ...(params.fastMode && + providerSupportsFastMode( + externalProvider.providerType, + externalSelection.modelId, + ) + ? { fast_mode: true } + : {}), ...(externalReasoningCaps.supportsReasoning ? externalReasoningCaps.reasoningStyle === "reasoning_effort" ? externalReasoningEnabled @@ -1603,6 +1654,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } continue; } + if (toolEvent.type === "anthropic_refusal") { + // Latch the backend refusal signal so the final + // message metadata can drive the prune. + anthropicRefusalSeen = true; + continue; + } if (toolEvent.type === "tool_start") { const id = (toolEvent.tool_call_id as string) || @@ -1936,6 +1993,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { timing: finalTiming, custom: { reasoningDuration, + // Persisted refusal flag driving the two-pass prune. + anthropicRefusal: anthropicRefusalSeen || undefined, serverTimings: meta?.timings ?? undefined, contextUsage: meta?.usage ? { diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 9cd4db2705..ac1ef8a24c 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -87,6 +87,7 @@ import { type ProviderCapabilities, getExternalMinOutputTokens, providerSupportsBuiltinCodeExecution, + providerSupportsFastMode, } from "./provider-capabilities"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import type { InferenceParams } from "./types/runtime"; @@ -552,6 +553,12 @@ export function ChatSettingsPanel({ activeExternalProvider.baseUrl, ) && activeExternalProvider.providerType === "openai"; + const showFastModeControl = + activeExternalProvider != null && + providerSupportsFastMode( + activeExternalProvider.providerType, + externalSelection?.modelId, + ); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const openAiApiKeyForSection = activeExternalProvider ? getExternalProviderApiKey(activeExternalProvider.id) || null @@ -1152,6 +1159,28 @@ export function ChatSettingsPanel({
) : null} + {showFastModeControl ? ( +
+
+ + Fast mode + + + Beta. Up to 2.5x higher output tokens per second on + Claude Opus 4.6 and 4.7 at 6x standard Opus pricing. + Switching between fast and standard invalidates the + prompt cache and is incompatible with the Priority + service tier. + +
+ +
+ ) : null} ) : null} diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index da1d6e3431..562a60a18f 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -135,6 +135,30 @@ export function providerSupportsBuiltinWebFetch( return providerType === "anthropic"; } +/** + * Whether the active provider + model supports Anthropic fast-mode + * (`speed: "fast"` + `fast-mode-2026-02-01` header). Opus 4.6 / 4.7 + * only per https://platform.claude.com/docs/en/build-with-claude/fast-mode. + * Backend silently drops on unsupported models as a second defence. + */ +const ANTHROPIC_FAST_MODE_MODEL_PREFIXES = [ + "claude-opus-4-7", + "claude-opus-4-6", +] as const; + +export function providerSupportsFastMode( + providerType: string | null | undefined, + modelId: string | null | undefined, +): boolean { + if (providerType !== "anthropic") return false; + if (!modelId) return false; + // Family boundary ("" or "-") required so IDs like "claude-opus-4-70" + // / "claude-opus-4-7b" do not match. + return ANTHROPIC_FAST_MODE_MODEL_PREFIXES.some( + (prefix) => modelId === prefix || modelId.startsWith(`${prefix}-`), + ); +} + /** * Whether the selected external provider/model exposes a server-side * code-execution tool. Two providers ship one today: 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 6b60ed51ea..73266b9234 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -379,6 +379,7 @@ const PERSISTED_INFERENCE_PARAM_KEYS = [ "maxTokens", "systemPrompt", "trustRemoteCode", + "fastMode", ] as const satisfies readonly PersistedInferenceParamKey[]; const SCALAR_SETTING_KEYS = [ diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 1e6bcf8b87..f18407413d 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -262,6 +262,12 @@ export interface OpenAIChatCompletionsRequest { * the Anthropic provider with `code_execution` in `enabled_tools`. */ anthropic_code_exec_container_id?: string | null; + /** + * Anthropic fast-mode toggle. Opus 4.6 / 4.7 only; backend drops + * silently on every other model + provider. See + * https://platform.claude.com/docs/en/build-with-claude/fast-mode + */ + fast_mode?: boolean | null; } export interface OpenAIChatDelta { diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 2967584653..4c44ee1e9c 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -14,6 +14,12 @@ export interface InferenceParams { checkpoint: string; /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trustRemoteCode?: boolean; + /** + * Anthropic fast-mode toggle. Opus 4.6 / 4.7 only; higher OTPS at + * 6x standard Opus pricing. Default false. + * https://platform.claude.com/docs/en/build-with-claude/fast-mode + */ + fastMode?: boolean; } export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { @@ -28,6 +34,7 @@ export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { systemPrompt: "", checkpoint: "", trustRemoteCode: false, + fastMode: false, }; export interface ChatModelSummary { diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index e07e1ddb1d..4e93a20bff 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -140,6 +140,11 @@ function sanitizeInferenceParams( if (typeof value.trustRemoteCode === "boolean") { params.trustRemoteCode = value.trustRemoteCode; } + // Mirror trustRemoteCode handling so the toggle survives reload + // and the /api/chat/settings round-trip. + if (typeof value.fastMode === "boolean") { + params.fastMode = value.fastMode; + } return hasKeys(params) ? params : undefined; } From 063e1e497b8922bada2854b042ce59d982c6af6a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 23:37:16 -0700 Subject: [PATCH 12/43] Studio: rewrite OpenAI Responses citation markers to markdown links (#5713) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: rewrite OpenAI Responses citation markers to markdown links OpenAI's /v1/responses stream interleaves text deltas with inline citation markers built from private-use codepoints (U+E200 / U+E201 / U+E202) shaped like `citeSOURCE_ID`. The codepoints render as garbled "E202" glyphs or empty boxes in most fonts, and the markdown layer further strips them, leaving run-on text like "citeturn1view0turn1view1turn3view0...". The url list still arrived in the Sources panel via url_citation annotations, but the inline cite hand-off into the prose was unreadable. Rewrite each marker into `[N](URL)` when the matching url_citation has already been recorded on this stream, and drop the marker silently otherwise. The lookup uses a new `source_id` field captured on `_record_url_citation` (accepts source_id / id / locator across Responses API revisions). Annotations are now applied BEFORE the delta text is rewritten so that markers and their resolving annotation arriving in the same SSE event still resolve. Reference: https://developers.openai.com/api/docs/guides/citation-formatting * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Preserve every source_id alias for a deduplicated url_citation OpenAI's Responses stream cites the same URL under multiple source_id markers when the model references different spans of the same page. The previous dedup-by-URL kept only the first alias and dropped the rest, so subsequent markers for the same URL never resolved and got stripped from the prose. Switch the citation record to a ``source_ids`` list and append new aliases on every duplicate. The rewriter resolves any alias back to the same citation number so the inline markers all collapse onto one footnote rather than fanning out into bogus repeats. Also collapse the two passes over ``all_url_citations`` in ``_record_url_citation`` into a single loop for clarity. Adds two regression tests covering the alias-collision and mixed-shape cases. * ci: re-trigger after flake in Studio GGUF Tool calling (rebased on main #5741 already) * ci: re-run after transient CodeQL Python checkout auth flake * Fix split-marker buffer + multi-source ids for PR #5713 The original rewriter only handles markers that arrive whole inside a single response.output_text.delta event. OpenAI's stream chunks text on byte-buffer boundaries with no awareness of the marker grammar, so a marker can straddle two deltas (delta-1 ends with "citetu", delta-2 starts with "rn0view0"). Each delta was rewritten in isolation, so the half-marker leaked as garbled "E200/E202" glyphs in the rendered prose. Buffer the unterminated tail across deltas and concatenate it onto the front of the next one so the rewriter sees a complete marker. Flush the held-over tail on response.completed / response.incomplete / [DONE], stripping any leftover private-use bytes so a never-closed marker (truncated stream, missing annotation) never leaks. Also handle the multi-source marker shape from the OpenAI docs -- citeid1id2 should expand to one bracket link per resolvable id. The previous regex captured only the first source id and silently dropped id2/id3. Reference: https://developers.openai.com/api/docs/guides/citation-formatting Tests: 21 new cases covering multi-source, locator suffix, marker split across two and three deltas, unterminated marker on truncation, late annotation resolving a buffered marker, idempotency, and the head/tail split helper directly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Defer citation segments until url_citation annotation arrives The split-marker buffer already concatenates a marker that straddles two response.output_text.delta events. But when the annotation event for a url_citation arrives AFTER the delta that contains its inline marker (the typical OpenAI Responses ordering), the rewriter still saw an empty lookup table at delta time and silently stripped the marker. The URL kept showing up in the sources panel but the inline link reference was permanently gone. Add _rewrite_citation_markers_partial which leaves an unresolved marker verbatim and reports has_unresolved=True. The streaming loop buffers any closed segment that contains an unresolved marker into a pending_citation_segments FIFO and drains the queue on every later annotation event, on response.completed, on response.incomplete, and on the [DONE] sentinel. Drain order is preserved so later clean text does not leapfrog an earlier deferred segment. End-of-stream forces a strip so no codepoint leaks if the annotation never arrived. Add six regression tests covering single-pass resolution, the late- annotation two-pass case, multi-source markers with partial resolution, mixed known and pending markers in one segment, and idempotency on marker-free input. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop unterminated citation tail to prevent cite-prefix plain-text leak `_flush_pending_marker_tail` stripped the three private-use citation codepoints from the held-over buffer, but left the literal ``cite`` keyword plus the source id behind as plain text. A stream ending mid-marker therefore emitted user-visible garbage like ``Some text citeturn0view0`` instead of the intended clean prose. ``pending_marker_tail`` is by construction the suffix that starts at an unclosed ``\\ue200`` opener -- the split helper guarantees there is no closing ``\\ue201`` byte. Without that close the marker is meaningless: the source id cannot be resolved to a URL and the user prose before the opener was already emitted as ``head`` on the originating delta. Bail out before the strip step and return the empty string. As a belt-and-braces measure also drop any orphan ``cite`` literal at the head of the buffer in case a future caller passes a partially-terminated tail. Update the matching ``_simulate_delta_stream`` harness in the edge tests so it mirrors the new flush logic, and add four regression tests covering unterminated marker with surrounding prose, marker- only inputs, prefix-only outputs, and the split-then-close path that still must resolve to a link. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Defer multi-source markers until all ids resolve for PR #5713 `_rewrite_citation_markers_partial` previously treated a marker as resolved when even one token in a multi-source marker resolved, dropping any still-pending source ids. In streamed Responses events the annotations for a multi-source marker can arrive across separate `annotation.added` chunks, so the caller no longer buffered that segment for retry and the late source id was lost from the inline citation entirely. Flag the marker unresolved whenever any token misses the lookup so the streamer keeps the segment pending. End-of-stream force flush still drops unresolved tokens through `_replace_openai_citation_markers` so locator-style suffixes (which look like unresolved ids at the token level but only appear at end-of-stream) render cleanly. Updated the multi-source test to assert the new pending-then-flush behavior; locator output now lands at force-flush rather than mid stream. * Shorten citation marker comments for PR #5713 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 343 ++++++++++++++- .../tests/test_openai_citation_markers.py | 251 +++++++++++ .../test_openai_citation_markers_edge.py | 413 ++++++++++++++++++ 3 files changed, 994 insertions(+), 13 deletions(-) create mode 100644 studio/backend/tests/test_openai_citation_markers.py create mode 100644 studio/backend/tests/test_openai_citation_markers_edge.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index d8a36610b8..0904426633 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -68,6 +68,136 @@ _ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile( ) _OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)") +# OpenAI Responses inline citation markers: `citeSOURCE_ID[id2...][LOCATOR]` +# using private-use codepoints (see +# https://developers.openai.com/api/docs/guides/citation-formatting). +# Group 1 holds the delim-separated tokens; each resolvable token expands +# to `[[N]](URL)`, unresolved tokens (locators, unknown ids) drop silently +# so no garbled glyph reaches the renderer. +_OPENAI_CITE_OPEN = "cite" +_OPENAI_CITE_STOP = "" +_OPENAI_CITE_DELIM = "" +_OPENAI_CITATION_MARKER = re.compile( + f"{_OPENAI_CITE_OPEN}([^{_OPENAI_CITE_STOP}]+){_OPENAI_CITE_STOP}" +) + + +def _build_citation_lookup( + url_citations: list[dict[str, Any]], +) -> dict[str, tuple[int, str]]: + """Map every known ``source_id`` alias to ``(citation_index, url)``. + + Accepts singular ``source_id`` and plural ``source_ids``. First-seen + wins on alias collision so an earlier citation keeps its number. + """ + by_source: dict[str, tuple[int, str]] = {} + for idx, cit in enumerate(url_citations, start = 1): + url = cit.get("url") + if not isinstance(url, str) or not url: + continue + aliases: list[str] = [] + sid = cit.get("source_id") + if isinstance(sid, str) and sid: + aliases.append(sid) + sids = cit.get("source_ids") + if isinstance(sids, list): + aliases.extend(s for s in sids if isinstance(s, str) and s) + for alias in aliases: + by_source.setdefault(alias, (idx, url)) + return by_source + + +def _replace_openai_citation_markers( + text: str, + url_citations: list[dict[str, Any]], +) -> str: + """Rewrite `\\ue200cite\\ue202SOURCE_ID[\\ue202LOCATOR]\\ue201` markers into + `[[N]](URL)` per resolvable id. Multi-source markers expand to one link + per id; unresolved tokens drop silently. Idempotent on text without + private-use codepoints. + """ + if not text or _OPENAI_CITE_STOP not in text: + return text + by_source = _build_citation_lookup(url_citations) + + def _sub(match: re.Match[str]) -> str: + # Try every delim-split token; unresolved tokens drop silently. + # Handles multi-source (all resolve) and source+locator (only the + # id resolves, locator drops). Empty result strips the marker. + rendered: list[str] = [] + for tok in match.group(1).split(_OPENAI_CITE_DELIM): + if not tok: + continue + hit = by_source.get(tok) + if hit is None: + continue + idx, url = hit + rendered.append(f"[[{idx}]]({url})") + return "".join(rendered) + + return _OPENAI_CITATION_MARKER.sub(_sub, text) + + +def _rewrite_citation_markers_partial( + text: str, + url_citations: list[dict[str, Any]], +) -> tuple[str, bool]: + """Like ``_replace_openai_citation_markers`` but also reports whether + any marker referenced a source_id not yet in ``url_citations``. + + The ``annotation.added`` event for a url_citation typically arrives + AFTER the delta carrying the marker referencing it. Callers buffer the + segment until a later event records the annotation; unresolved markers + are left verbatim so a follow-up pass still parses cleanly. + """ + if not text or _OPENAI_CITE_STOP not in text: + return text, False + by_source = _build_citation_lookup(url_citations) + has_unresolved = False + + def _sub(match: re.Match[str]) -> str: + nonlocal has_unresolved + tokens = [t for t in match.group(1).split(_OPENAI_CITE_DELIM) if t] + rendered: list[str] = [] + any_unresolved = False + for tok in tokens: + hit = by_source.get(tok) + if hit is None: + any_unresolved = True + continue + idx, url = hit + rendered.append(f"[[{idx}]]({url})") + # Leave the whole marker verbatim if any token is unresolved so the + # caller can re-run once the late annotation lands; partial emission + # would lose the unresolved ids once the source text is dropped. + if any_unresolved: + has_unresolved = True + return match.group(0) + return "".join(rendered) + + return _OPENAI_CITATION_MARKER.sub(_sub, text), has_unresolved + + +def _split_pending_citation_tail(text: str) -> tuple[str, str]: + """Split ``text`` into ``(head, pending_tail)`` for streamed deltas. + + A citation marker can straddle two SSE deltas (e.g. delta-1 ends with + ``\\ue200citetu`` and delta-2 starts with ``rn0view0\\ue201``); the + unterminated tail is buffered and prepended onto the next delta so the + rewriter sees a complete marker. ``pending_tail`` is the longest suffix + starting with ``\\ue200`` and lacking ``\\ue201``; ``head`` is safe to + emit. Empty tail when ``text`` has no open marker or a fully closed one. + """ + if not text: + return text, "" + last_open = text.rfind("") + if last_open == -1: + return text, "" + # Stop byte after the last open byte means the marker closed in this delta. + if _OPENAI_CITE_STOP in text[last_open:]: + return text, "" + return text[:last_open], text[last_open:] + class _AnthropicThinkingSpec(NamedTuple): prefixes: tuple[str, ...] @@ -3000,6 +3130,65 @@ class ExternalProviderClient: # see. latched_container_id: Optional[str] = None container_id_emitted = False + # Buffer for a citation marker straddling two delta events; + # prepended onto the next delta. See _split_pending_citation_tail. + pending_marker_tail: str = "" + # Segments deferred while their markers reference unseen + # source_ids; held in arrival order so output never + # leapfrogs an earlier deferred segment. Flushed on + # annotation events and force-flushed at end-of-stream + # with leftover private-use codepoints stripped. + pending_citation_segments: list[str] = [] + + def _drain_pending_segments(force: bool) -> str: + """Re-attempt resolution on buffered segments in order. + Stops at the first still-unresolved segment unless + ``force`` (end-of-stream), where lingering markers are stripped.""" + out: list[str] = [] + while pending_citation_segments: + seg = pending_citation_segments[0] + rewritten, unresolved = _rewrite_citation_markers_partial( + seg, + all_url_citations, + ) + if unresolved and not force: + pending_citation_segments[0] = rewritten + break + if unresolved and force: + rewritten = _replace_openai_citation_markers( + rewritten, + all_url_citations, + ) + pending_citation_segments.pop(0) + if rewritten: + out.append(rewritten) + return "".join(out) + + def _flush_pending_marker_tail(tail: str) -> str: + """Render any leftover citation tail at end-of-stream. + + Unterminated tails drop (no annotation to bind to). If the + close byte arrived concatenated, rewrite then scrub any + residual private-use bytes and any orphan ``cite`` + literal so the renderer never sees raw markup. url_citations + are aggregated separately and applied to web_search tool_end. + """ + if not tail: + return "" + if _OPENAI_CITE_STOP not in tail: + # Unterminated: drop the whole tail, otherwise the + # residual ``cite`` would leak as plain text. + return "" + rendered = _replace_openai_citation_markers( + tail, all_url_citations + ) + # Scrub residual private-use bytes (e.g. a partial opener). + for ch in ("", "", ""): + rendered = rendered.replace(ch, "") + # Drop any orphan ``cite`` literal -- meaningless + # without its closing byte and matching url_citation. + rendered = re.sub(r"^cite\S*", "", rendered) + return rendered def _emit_tool_event(payload: dict[str, Any]) -> str: chunk = { @@ -3057,16 +3246,35 @@ class ExternalProviderClient: def _record_url_citation(payload: dict[str, Any]) -> None: """Append a url_citation onto the shared all_url_citations - list. Dedup by URL — the same source can be cited multiple - times across deltas. We do NOT try to attribute citations - to individual web_search_call invocations because OpenAI's - annotation events don't carry that linkage.""" + list. Dedup by URL — the same URL can be cited many + times under different ``source_id`` aliases (one per + span/locator), so collect every alias we see onto + the matching entry's ``source_ids`` list. The + delta-text rewriter resolves any of those aliases + back to this entry's URL. The id may live under + ``source_id``, ``id``, or ``locator`` across the + Responses API revisions.""" if payload.get("type") != "url_citation": return url = payload.get("url", "") if not url: return - if any(c["url"] == url for c in all_url_citations): + source_id = ( + payload.get("source_id") + or payload.get("id") + or payload.get("locator") + or "" + ) + # Single pass: either backfill aliases onto an + # existing URL entry (and return) or fall through + # to append a fresh one. + for c in all_url_citations: + if c["url"] != url: + continue + if source_id: + aliases = c.setdefault("source_ids", []) + if source_id not in aliases: + aliases.append(source_id) return title = payload.get("title") or url snippet = payload.get("snippet") or payload.get("quote") or "" @@ -3075,6 +3283,7 @@ class ExternalProviderClient: "url": url, "title": title, "snippet": snippet, + "source_ids": [source_id] if source_id else [], } ) @@ -3131,6 +3340,28 @@ class ExternalProviderClient: if not data_str: continue if data_str == "[DONE]": + # Flush any held-over partial marker; strip + # private-use bytes so garbled glyphs don't leak. + if pending_marker_tail: + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) + pending_marker_tail = "" + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) + # Force-drain any segment still awaiting an + # annotation; lingering codepoints are stripped. + tail_flushed = _drain_pending_segments( + force = True, + ) + if tail_flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(tail_flushed) if not done_emitted: yield "data: [DONE]" done_emitted = True @@ -3145,22 +3376,57 @@ class ExternalProviderClient: if event_type == "response.output_text.delta": delta_text = event.get("delta", "") - if delta_text: - if reasoning_open: - yield _chunk_with_text("") - reasoning_open = False - yield _chunk_with_text(delta_text) - # Some API versions inline url citations on the - # delta event itself rather than as a separate - # response.output_text.annotation.added event. + # Process inline annotations first so source_ids + # referenced by same-delta markers are in the lookup + # before the rewriter runs. Some API versions inline + # url citations on the delta event itself. for ann in event.get("annotations") or []: if isinstance(ann, dict): _record_url_citation(ann) + if delta_text or pending_marker_tail: + # Prepend any held-over tail so a marker + # straddling two SSE events resolves cleanly. + combined = pending_marker_tail + delta_text + head, pending_marker_tail = ( + _split_pending_citation_tail(combined) + ) + if head: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + # Re-attempt earlier deferred segments first + # so output stays in order; the needed + # annotation may have arrived inline above. + flushed = _drain_pending_segments( + force = False, + ) + if flushed: + yield _chunk_with_text(flushed) + head_rewritten, has_unresolved = ( + _rewrite_citation_markers_partial( + head, + all_url_citations, + ) + ) + if has_unresolved or pending_citation_segments: + pending_citation_segments.append( + head_rewritten + ) + elif head_rewritten: + yield _chunk_with_text(head_rewritten) elif event_type == "response.output_text.annotation.added": ann = event.get("annotation") if isinstance(ann, dict): _record_url_citation(ann) + flushed = _drain_pending_segments( + force = False, + ) + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) elif event_type == "response.output_item.added": # Track the call early but do NOT emit tool_start @@ -3399,6 +3665,34 @@ class ExternalProviderClient: ) if isinstance(completed_usage, dict): last_usage = completed_usage + # Flush any unterminated citation tail + # held over from the last delta. By + # the time we get here every annotation + # has been recorded so a late-arriving + # source_id may resolve cleanly; if it + # still doesn't, the helper strips the + # private-use bytes so no garbled + # glyph reaches the user. + if pending_marker_tail: + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) + pending_marker_tail = "" + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) + # Force-drain any segment still awaiting an + # annotation; lingering codepoints are stripped. + tail_flushed = _drain_pending_segments( + force = True, + ) + if tail_flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(tail_flushed) if reasoning_open: yield _chunk_with_text("") reasoning_open = False @@ -3491,6 +3785,29 @@ class ExternalProviderClient: ) if isinstance(incomplete_usage, dict): last_usage = incomplete_usage + # Same flush as response.completed -- + # truncated streams can leave a half- + # marker in the buffer. + if pending_marker_tail: + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) + pending_marker_tail = "" + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) + # Force-drain any segment still awaiting an + # annotation; lingering codepoints are stripped. + tail_flushed = _drain_pending_segments( + force = True, + ) + if tail_flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(tail_flushed) if reasoning_open: yield _chunk_with_text("") reasoning_open = False diff --git a/studio/backend/tests/test_openai_citation_markers.py b/studio/backend/tests/test_openai_citation_markers.py new file mode 100644 index 0000000000..ccc17be329 --- /dev/null +++ b/studio/backend/tests/test_openai_citation_markers.py @@ -0,0 +1,251 @@ +# 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 OpenAI Responses-API citation marker rewriter. + +The stream interleaves text deltas with ``\\ue200cite\\ue202SOURCE_ID\\ue201`` +markers. The rewriter resolves each to `[N](URL)` when the annotation has +arrived and drops it otherwise; the URL list still flows to Sources via +`_record_url_citation`. + +Reference: https://developers.openai.com/api/docs/guides/citation-formatting +""" + +import pytest + +from core.inference.external_provider import ( + _replace_openai_citation_markers, + _rewrite_citation_markers_partial, +) + + +# Citation marker control codepoints (private-use area): +CITE_START = "" +CITE_STOP = "" +CITE_DELIM = "" + + +def _marker(source_id: str, locator: str | None = None) -> str: + payload = f"{CITE_START}cite{CITE_DELIM}{source_id}" + if locator: + payload = f"{payload}{CITE_DELIM}{locator}" + return f"{payload}{CITE_STOP}" + + +def _has_marker_codepoints(text: str) -> bool: + return any(c in text for c in (CITE_START, CITE_STOP, CITE_DELIM)) + + +def test_passthrough_when_no_marker_present(): + text = "Plain text with no citation markers." + assert _replace_openai_citation_markers(text, []) == text + + +def test_marker_rewritten_to_link_when_annotation_known(): + text = f"The capital is Paris {_marker('turn0view0')}." + citations = [ + { + "source_id": "turn0view0", + "url": "https://example.com/paris", + "title": "Paris", + }, + ] + out = _replace_openai_citation_markers(text, citations) + assert not _has_marker_codepoints(out) + assert "[[1]](https://example.com/paris)" in out + + +def test_unknown_source_marker_dropped_silently(): + text = f"Foo {_marker('turn9view9')} bar." + out = _replace_openai_citation_markers(text, []) + # Marker stripped, no garbled "E202" glyph leaks through, and the + # surrounding text stays intact. + assert not _has_marker_codepoints(out) + assert "E202" not in out + assert "turn9view9" not in out + assert "Foo" in out and "bar" in out + + +def test_multiple_concatenated_markers_resolved_in_order(): + """Real-world wire shape: a string of markers butted up against each other + after a sentence, as in the user-reported bug.""" + markers = "".join(_marker(f"turn{i}view{j}") for i, j in [(1, 0), (1, 1), (3, 0)]) + text = f"All animals ranked. {markers}" + citations = [ + {"source_id": "turn1view0", "url": "https://a.example/dog", "title": "Dog"}, + {"source_id": "turn1view1", "url": "https://a.example/cat", "title": "Cat"}, + {"source_id": "turn3view0", "url": "https://a.example/tiger", "title": "Tiger"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://a.example/dog)" in out + assert "[[2]](https://a.example/cat)" in out + assert "[[3]](https://a.example/tiger)" in out + assert not _has_marker_codepoints(out) + + +def test_marker_with_locator_resolves(): + text = f"See {_marker('turn2file0', 'L8-L13')}." + citations = [ + {"source_id": "turn2file0", "url": "https://example.com/doc.txt"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/doc.txt)" in out + assert "L8-L13" not in out # locator detail dropped; we just link. + assert not _has_marker_codepoints(out) + + +def test_mixed_known_and_unknown_markers(): + known = _marker("turn0view0") + unknown = _marker("turn0view99") + text = f"Known {known} and unknown {unknown}." + citations = [ + {"source_id": "turn0view0", "url": "https://example.com/known"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/known)" in out + # Unknown markers leave no trace, but surrounding prose stays. + assert "Known" in out and "unknown" in out + assert not _has_marker_codepoints(out) + assert "E202" not in out + + +def test_empty_text_returns_verbatim(): + assert _replace_openai_citation_markers("", []) == "" + + +def test_idempotent_on_pre_stripped_text(): + """Pre-stripped text (no private-use codepoints) returns verbatim.""" + text = "citeturn1view0 plain" + assert _replace_openai_citation_markers(text, []) == text + + +@pytest.mark.parametrize( + "citation", + [ + {"url": "https://example.com/a"}, # no source_id at all + {"source_id": None, "url": "https://example.com/b"}, + {"source_id": "", "url": "https://example.com/c"}, + ], +) +def test_citation_without_source_id_does_not_crash(citation): + text = f"X {_marker('turnXviewY')} Y" + out = _replace_openai_citation_markers(text, [citation]) + # No mapping, marker stripped. Crash-free is the contract. + assert not _has_marker_codepoints(out) + assert "turnXviewY" not in out + + +def test_multiple_source_id_aliases_resolve_to_same_url(): + """Every alias for the same URL must resolve, not just the first. + Regression for the Codex P1 on the original PR.""" + a = _marker("turn0view0") + b = _marker("turn0view0_span_1") + c = _marker("turn0view0_span_2") + text = f"Triple {a}{b}{c} cite." + citations = [ + { + "source_ids": ["turn0view0", "turn0view0_span_1", "turn0view0_span_2"], + "url": "https://example.com/paris", + "title": "Paris", + }, + ] + out = _replace_openai_citation_markers(text, citations) + # All three aliases collapse onto citation [1] -- the URL is the + # same so it would be misleading to show three different numbers. + assert out.count("[[1]](https://example.com/paris)") == 3 + assert not _has_marker_codepoints(out) + + +def test_source_ids_list_and_legacy_source_id_both_resolve(): + """Mixed-shape citation: legacy ``source_id`` plus newer + ``source_ids`` aliases both resolve.""" + legacy = _marker("legacy_id") + alias = _marker("alias_id") + text = f"Both {legacy} and {alias} work." + citations = [ + { + "source_id": "legacy_id", + "source_ids": ["alias_id"], + "url": "https://example.com/doc", + }, + ] + out = _replace_openai_citation_markers(text, citations) + assert out.count("[[1]](https://example.com/doc)") == 2 + assert not _has_marker_codepoints(out) + + +# --------------------------------------------------------------------------- +# _rewrite_citation_markers_partial: deferred-annotation tests. OpenAI emits +# url_citation annotations on a subsequent SSE event; this helper reports +# `has_unresolved` so the stream loop defers emission. See PR #5713 audit. +# --------------------------------------------------------------------------- + + +def test_partial_known_marker_resolves_and_clears_unresolved(): + text = f"Foo {_marker('s1')} bar." + out, unresolved = _rewrite_citation_markers_partial( + text, + [{"source_id": "s1", "url": "https://example.com/a"}], + ) + assert "[[1]](https://example.com/a)" in out + assert unresolved is False + assert not _has_marker_codepoints(out) + + +def test_partial_unknown_marker_preserves_verbatim_and_flags(): + text = f"Foo {_marker('s1')} bar." + out, unresolved = _rewrite_citation_markers_partial(text, []) + assert unresolved is True + # Codepoints must remain so a follow-up pass can re-parse. + assert _has_marker_codepoints(out) + assert "Foo" in out and "bar." in out + + +def test_partial_resolves_after_late_annotation(): + """Two-pass: first call sees no citations, second resolves after annotation.""" + text = f"See {_marker('s1')} for details." + out1, unresolved1 = _rewrite_citation_markers_partial(text, []) + assert unresolved1 is True + citations = [{"source_id": "s1", "url": "https://example.com/x"}] + out2, unresolved2 = _rewrite_citation_markers_partial(out1, citations) + assert unresolved2 is False + assert "[[1]](https://example.com/x)" in out2 + assert not _has_marker_codepoints(out2) + + +def test_partial_multi_source_partial_resolution_keeps_marker_pending(): + """Any unresolved token in a multi-source marker leaves the whole marker + verbatim with ``unresolved`` True; defer until every id resolves or + end-of-stream forces a flush (dropping unresolved tokens then).""" + cite = f"{CITE_START}cite{CITE_DELIM}known{CITE_DELIM}locator{CITE_STOP}" + text = f"Pre {cite} post." + citations = [{"source_id": "known", "url": "https://example.com/y"}] + out, unresolved = _rewrite_citation_markers_partial(text, citations) + assert unresolved is True + assert cite in out + # End-of-stream force flush: drop the unresolved token, keep the + # resolved link. The streamer routes pending segments through + # `_replace_openai_citation_markers` at force=True for this. + forced = _replace_openai_citation_markers(out, citations) + assert "[[1]](https://example.com/y)" in forced + assert "locator" not in forced + assert not _has_marker_codepoints(forced) + + +def test_partial_idempotent_on_marker_free_text(): + text = "Plain text." + out, unresolved = _rewrite_citation_markers_partial(text, []) + assert out == text + assert unresolved is False + + +def test_partial_mixed_known_and_pending_markers_flags_unresolved(): + known = _marker("known") + pending = _marker("pending") + text = f"{known} {pending}" + citations = [{"source_id": "known", "url": "https://example.com/k"}] + out, unresolved = _rewrite_citation_markers_partial(text, citations) + assert unresolved is True # the pending marker drives the flag + assert "[[1]](https://example.com/k)" in out + # The pending marker stays verbatim for the next pass. + assert CITE_START in out and "pending" in out diff --git a/studio/backend/tests/test_openai_citation_markers_edge.py b/studio/backend/tests/test_openai_citation_markers_edge.py new file mode 100644 index 0000000000..ffe8c6b6eb --- /dev/null +++ b/studio/backend/tests/test_openai_citation_markers_edge.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case tests for the OpenAI Responses citation marker rewriter. + +Covers multi-source markers, source+locator, marker SPLIT across SSE deltas, +unterminated tails at end-of-stream, multiple markers per delta, late +annotation ordering, and idempotency. + +Reference: https://developers.openai.com/api/docs/guides/citation-formatting +""" + +import importlib + + +# Streaming integration is exercised by ``_simulate_delta_stream`` further +# down, mirroring the head/buffer/flush dance from ``_stream_openai_responses``. +_module = importlib.import_module("core.inference.external_provider") +_replace_openai_citation_markers = _module._replace_openai_citation_markers +_split_pending_citation_tail = _module._split_pending_citation_tail + + +CITE_START = "" +CITE_STOP = "" +CITE_DELIM = "" + + +def _marker(*source_ids: str, locator: str | None = None) -> str: + """Build a ``\\ue200cite\\ue202[\\ue202...][\\ue202]\\ue201`` + marker. Accepts one or many ``source_ids`` plus an optional ``locator``.""" + payload = f"{CITE_START}cite{CITE_DELIM}" + CITE_DELIM.join(source_ids) + if locator: + payload = f"{payload}{CITE_DELIM}{locator}" + return f"{payload}{CITE_STOP}" + + +def _no_private_use(text: str) -> bool: + return all(c not in text for c in (CITE_START, CITE_STOP, CITE_DELIM)) + + +# Harness mirroring the head/pending-tail/flush dance in +# `_stream_openai_responses`, so streaming tests skip the httpx mock. +def _simulate_delta_stream( + deltas: list[str], + citations: list[dict], + *, + flush: bool = True, +) -> str: + pending = "" + emitted: list[str] = [] + for delta in deltas: + combined = pending + delta + head, pending = _split_pending_citation_tail(combined) + if head: + head = _replace_openai_citation_markers(head, citations) + if head: + emitted.append(head) + if flush and pending: + # Mirror `_flush_pending_marker_tail`: drop the tail entirely if no + # closing stop byte arrived; the literal ``cite`` would leak otherwise. + if CITE_STOP not in pending: + rendered = "" + else: + rendered = _replace_openai_citation_markers(pending, citations) + for ch in (CITE_START, CITE_STOP, CITE_DELIM): + rendered = rendered.replace(ch, "") + import re as _re + + rendered = _re.sub(r"^cite\S*", "", rendered) + if rendered: + emitted.append(rendered) + return "".join(emitted) + + +# --------------------------------------------------------------------------- +# 1. Multi-source markers per the OpenAI docs. +# --------------------------------------------------------------------------- + + +def test_multi_source_marker_all_resolve(): + """\\ue200cite\\ue202id1\\ue202id2\\ue202id3\\ue201 expands to three links + when every id is known. Earlier regex captured only id1 and dropped id2/id3.""" + text = f"All three: {_marker('id1', 'id2', 'id3')}" + citations = [ + {"source_id": "id1", "url": "https://example.com/1"}, + {"source_id": "id2", "url": "https://example.com/2"}, + {"source_id": "id3", "url": "https://example.com/3"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/1)" in out + assert "[[2]](https://example.com/2)" in out + assert "[[3]](https://example.com/3)" in out + assert _no_private_use(out) + + +def test_multi_source_marker_partial_resolution(): + """Known ids render, unknown ids drop silently, no glyph leaks.""" + text = f"Mixed: {_marker('known', 'unknown', 'also_known')}" + citations = [ + {"source_id": "known", "url": "https://k.example"}, + {"source_id": "also_known", "url": "https://ak.example"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://k.example)" in out + assert "[[2]](https://ak.example)" in out + assert "unknown" not in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 2. Source + locator: locator is dropped, link still resolves. +# --------------------------------------------------------------------------- + + +def test_marker_with_numeric_locator(): + text = f"See {_marker('tu0', locator = '42')}." + citations = [{"source_id": "tu0", "url": "https://example.com/doc"}] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/doc)" in out + assert "42" not in out + assert _no_private_use(out) + + +def test_marker_with_range_locator(): + text = f"See {_marker('tu0', locator = 'L8-L13')}." + citations = [{"source_id": "tu0", "url": "https://example.com/code"}] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/code)" in out + assert "L8-L13" not in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 3. Marker SPLIT across two SSE deltas -- the codex-flagged P1. +# --------------------------------------------------------------------------- + + +def test_marker_split_in_source_id(): + """Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 starts + with the rest (``rn0view0\\ue201``). The buffer stitches the halves + back together so they resolve to one link instead of leaking.""" + full = f"See {_marker('turn0view0')} now." + # Cut right after the second delim + "tu" inside the source id. + cut = full.index("tu", full.index(CITE_START)) + len("tu") + d1, d2 = full[:cut], full[cut:] + # Sanity check: delta-1 actually contains a partial marker. + assert CITE_START in d1 and CITE_STOP not in d1 + assert CITE_STOP in d2 + citations = [{"source_id": "turn0view0", "url": "https://x"}] + out = _simulate_delta_stream([d1, d2], citations) + assert out == "See [[1]](https://x) now." + assert _no_private_use(out) + + +def test_marker_split_at_start_byte(): + """Split exactly after the opening ``\\ue200`` byte; the buffer must + hold the lone open byte until the rest arrives.""" + full = f"Text {_marker('sid')} done" + cut = full.index(CITE_START) + 1 # right AFTER the open byte + d1, d2 = full[:cut], full[cut:] + citations = [{"source_id": "sid", "url": "https://y"}] + out = _simulate_delta_stream([d1, d2], citations) + assert out == "Text [[1]](https://y) done" + assert _no_private_use(out) + + +def test_marker_split_across_three_deltas(): + """Worst case: marker chopped into three pieces across three deltas.""" + full = f"A {_marker('threesplit')} B" + # cut at two points inside the marker + open_pos = full.index(CITE_START) + stop_pos = full.index(CITE_STOP) + cut1 = open_pos + 4 + cut2 = stop_pos - 2 + parts = [full[:cut1], full[cut1:cut2], full[cut2:]] + citations = [{"source_id": "threesplit", "url": "https://z"}] + out = _simulate_delta_stream(parts, citations) + assert out == "A [[1]](https://z) B" + assert _no_private_use(out) + + +def test_marker_split_with_trailing_text_after_close(): + """Delta-2 closes the marker AND carries trailing prose; both emit cleanly.""" + full = f"X {_marker('sid')} after" + cut = full.index("cite") + len("ci") + d1, d2 = full[:cut], full[cut:] + citations = [{"source_id": "sid", "url": "https://a"}] + out = _simulate_delta_stream([d1, d2], citations) + assert out == "X [[1]](https://a) after" + assert _no_private_use(out) + + +def test_split_marker_unknown_source_is_dropped_cleanly(): + """Split marker for an unknown source drops silently on flush.""" + full = f"Pre {_marker('never_seen')} post" + cut = full.index(CITE_START) + 3 + d1, d2 = full[:cut], full[cut:] + out = _simulate_delta_stream([d1, d2], []) + assert out == "Pre post" + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 4. Unterminated marker at end-of-stream -- truncation safety. +# --------------------------------------------------------------------------- + + +def test_unterminated_marker_at_stream_end_dropped_on_flush(): + """Stream ends mid-marker (e.g. response.incomplete); the tail is + flushed with private-use bytes stripped, no `E202` text leaks.""" + deltas = ["Some text ", f"{CITE_START}citetu", "rn0view0"] # no STOP ever + out = _simulate_delta_stream(deltas, [], flush = True) + assert _no_private_use(out) + assert "E200" not in out and "E202" not in out + # Surrounding prose stays; we don't assert exact marker remainder. + assert "Some text " in out + + +def test_flush_resolves_marker_when_late_annotation_arrives(): + """Marker in a delta, matching annotation arrives later (on + response.output_text.annotation.added after the final delta). The + rewriter reads ``all_url_citations`` LIVE at flush, so the buffered + marker still resolves.""" + deltas = ["Look ", f"{CITE_START}cite{CITE_DELIM}late_sid"] + pending = "" + citations: list[dict] = [] + emitted: list[str] = [] + for d in deltas: + combined = pending + d + head, pending = _split_pending_citation_tail(combined) + if head: + emitted.append(_replace_openai_citation_markers(head, citations)) + # Annotation arrives AFTER all deltas but BEFORE flush. + citations.append({"source_id": "late_sid", "url": "https://late.example"}) + # Append the STOP byte that closed the marker in a later delta. + pending = pending + CITE_STOP + flushed = _replace_openai_citation_markers(pending, citations) + for ch in (CITE_START, CITE_STOP, CITE_DELIM): + flushed = flushed.replace(ch, "") + emitted.append(flushed) + out = "".join(emitted) + assert "[[1]](https://late.example)" in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 5. Multiple unrelated markers in a single delta. +# --------------------------------------------------------------------------- + + +def test_three_markers_in_one_delta_resolve_independently(): + text = f"alpha {_marker('a')} beta {_marker('b')} gamma {_marker('c')} end" + citations = [ + {"source_id": "a", "url": "https://example.com/a"}, + {"source_id": "b", "url": "https://example.com/b"}, + {"source_id": "c", "url": "https://example.com/c"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert out == ( + "alpha [[1]](https://example.com/a) beta " + "[[2]](https://example.com/b) gamma " + "[[3]](https://example.com/c) end" + ) + + +# --------------------------------------------------------------------------- +# 6. Idempotency. +# --------------------------------------------------------------------------- + + +def test_rewriter_idempotent_on_already_rewritten_text(): + """Running the rewriter twice does not double-link or corrupt brackets.""" + text = f"alpha {_marker('a')} omega" + citations = [{"source_id": "a", "url": "https://example.com/a"}] + once = _replace_openai_citation_markers(text, citations) + twice = _replace_openai_citation_markers(once, citations) + assert once == twice + assert _no_private_use(once) + + +def test_rewriter_idempotent_on_marker_free_text(): + """No-op when there is nothing to rewrite.""" + text = "Plain prose with no citations and no private-use bytes." + out = _replace_openai_citation_markers(text, []) + assert out is text or out == text + + +# --------------------------------------------------------------------------- +# 7. Edge / robustness. +# --------------------------------------------------------------------------- + + +def test_only_marker_no_surrounding_text(): + """A delta that is JUST a marker (no prose) still renders correctly; + used to leak without the empty-string short-circuit in the split helper.""" + text = _marker("solo") + citations = [{"source_id": "solo", "url": "https://solo.example"}] + out = _replace_openai_citation_markers(text, citations) + assert out == "[[1]](https://solo.example)" + + +def test_back_to_back_markers_with_no_separator(): + """Adjacent markers resolve to concatenated links, no joining whitespace.""" + text = f"{_marker('x')}{_marker('y')}" + citations = [ + {"source_id": "x", "url": "https://x.example"}, + {"source_id": "y", "url": "https://y.example"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert out == "[[1]](https://x.example)[[2]](https://y.example)" + + +def test_split_helper_buffers_only_after_last_open_byte(): + """A complete marker followed by an unterminated one: head includes + the complete marker, buffer holds only the trailing partial.""" + complete = _marker("done") + partial = f"{CITE_START}cite{CITE_DELIM}half" # no STOP + text = f"pre {complete} mid {partial}" + head, tail = _split_pending_citation_tail(text) + assert head == f"pre {complete} mid " + assert tail == partial + # And the head, once rewritten, drops every private-use byte. + rewritten = _replace_openai_citation_markers( + head, [{"source_id": "done", "url": "https://d"}] + ) + assert rewritten == "pre [[1]](https://d) mid " + + +def test_split_helper_empty_input(): + head, tail = _split_pending_citation_tail("") + assert head == "" and tail == "" + + +def test_split_helper_no_open_byte(): + head, tail = _split_pending_citation_tail("nothing to see here") + assert head == "nothing to see here" and tail == "" + + +def test_split_helper_complete_marker_only(): + """A delta ending with a closed marker leaves the buffer empty.""" + text = f"alpha {_marker('a')}" + head, tail = _split_pending_citation_tail(text) + assert head == text and tail == "" + + +# --------------------------------------------------------------------------- +# 8. Sources-panel: marker drop must not affect citation aggregation. +# Indices come from the url_citations list, not the marker stream. +# --------------------------------------------------------------------------- + + +def test_unknown_marker_does_not_perturb_citation_indexing(): + """Unknown source_id markers drop without consuming an index slot.""" + text = f"A {_marker('unknown')} B {_marker('real_a')} C {_marker('real_b')}" + citations = [ + {"source_id": "real_a", "url": "https://example.com/a"}, + {"source_id": "real_b", "url": "https://example.com/b"}, + ] + out = _replace_openai_citation_markers(text, citations) + # real_a is index 1; unknown does not take a slot. + assert "[[1]](https://example.com/a)" in out + assert "[[2]](https://example.com/b)" in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# Regression: unterminated marker tail must NOT leak the residual +# ``cite``-prefixed source id as plain text. PR #5713 audit P1. +# --------------------------------------------------------------------------- + + +def test_unterminated_marker_does_not_leak_cite_residue(): + """Stream ends mid-marker: drop the whole tail rather than strip + codepoints and leave ``cite`` behind.""" + half = f"Hi there {CITE_START}cite{CITE_DELIM}turn0view0" + out = _simulate_delta_stream([half], [], flush = True) + # Prose before the marker stays; no private-use bytes or cite residue. + assert "Hi there" in out + assert _no_private_use(out) + assert "citeturn0view0" not in out + assert "cite" not in out.split("Hi there", 1)[1] + + +def test_unterminated_marker_only_no_prefix_drops_entirely(): + """A delta that is purely an unterminated marker flushes to "".""" + half = f"{CITE_START}cite{CITE_DELIM}turn0view0" + out = _simulate_delta_stream([half], [], flush = True) + assert out == "" + + +def test_unterminated_marker_with_prefix_emits_only_prefix(): + """Prose then unterminated marker: prose emits, marker remnant drops.""" + half = f"prefix prose {CITE_START}cite{CITE_DELIM}abc" + out = _simulate_delta_stream([half], [], flush = True) + assert out == "prefix prose " + + +def test_closing_byte_arrives_after_pending_buffered_split(): + """Closing byte arrives in a later delta after opener + source id were + buffered; link resolves with no residue.""" + cuts = [ + f"a {CITE_START}cite{CITE_DELIM}", + f"sid{CITE_STOP} b", + ] + out = _simulate_delta_stream( + cuts, + [{"source_id": "sid", "url": "https://example.com/x"}], + flush = True, + ) + assert "[[1]](https://example.com/x)" in out + assert "a " in out and "b" in out + assert _no_private_use(out) + assert "citesid" not in out From 7d3c472461cddf0530691eb7569676097b3868dd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 23:37:26 -0700 Subject: [PATCH 13/43] Studio: standalone Fetch pill for Anthropic web_fetch (#5742) * Studio: surface Anthropic web_fetch as a standalone Fetch pill web_fetch used to be silently bundled with the Search pill on the assumption that "search returns URLs, fetch reads them" is the typical workflow. Two problems with that: - Anthropic bills each web_fetch invocation separately from web_search hits, so combining them made the per-message cost surface ambiguous. - It blocked "just fetch this one URL" workflows where the user already knows the page they want read and does not want a search round-trip. Adds: - `webFetchToolsEnabled` to the chat-runtime-store, persisted to localStorage under `unsloth_chat_web_fetch_tools_enabled`, with a matching `supportsBuiltinWebFetch` capability flag and a `setWebFetchToolsEnabled` setter. - A new Fetch pill in the chat composer, rendered next to Images and only when the active provider returns true from `providerSupportsBuiltinWebFetch` (Anthropic today). The pill defaults off so per-fetch billing is always a deliberate opt-in. - chat-page bootstraps `webFetchToolsEnabled` from the same stored- preference fallback the other pills use. - chat-adapter reads `webFetchToolsEnabled` directly when deciding whether to append "web_fetch" to `enabled_tools`, decoupling it from `toolsEnabled` (Search). Backend translation is unchanged: when `enabled_tools` already contains "web_fetch", `_stream_anthropic` appends the `web_fetch_20250910` / `web_fetch_20260209` tool exactly as before (test_anthropic_web_fetch.py pins the standalone-only path at `test_web_fetch_tool_appended_to_request_body` and the combined path at `test_web_fetch_combined_with_web_search_and_code_execution`). Frontend tsc passes. * ci: re-trigger after transient GitHub API HTTP flake (checkout + ggml-org release fetch) * Studio: include web_fetch in the disabled-tool guard axis Reviewer P1 / High on PR #5742 (codex + gemini): after introducing the standalone Fetch pill, `disabledToolGuard` still only branched on `webSearchEnabledForThisTurn`. With Fetch ON and Search OFF the system prompt would tell Claude "you do not have web search or web fetch tools in this conversation", which contradicts the actual tool schema being sent and suppresses `web_fetch` tool calls, defeating the standalone-fetch workflow this PR adds. Treat search and fetch as a single "any web tool enabled" axis. The guard only needs to warn the model when no web tool is wired in for this turn; once either pill is on the model can pick the right one from the tool schema. The existing `webLabel` already covers both names, so the user-visible guard text stays accurate in every combination. tsc clean. * ci: re-trigger after transient infra flake on Windows prebuilt / actions/checkout * Studio: route web_fetch through per-model version dispatch The web_fetch tool body in `_stream_anthropic` hardcoded `web_fetch_20250910` instead of calling `_anthropic_web_fetch_version`, so Opus 4.6 / 4.7 and Sonnet 4.6 missed the `web_fetch_20260209` dynamic-filtering variant. The picker, the unit tests for it, and a deliberate "follow-up" note in `test_anthropic_web_fetch.py` already existed; this just threads it through the emission site. Mirrors how web_search and code_execution are dispatched per model. Old models still resolve to `web_fetch_20250910` and continue to work. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shorten web_fetch comments for PR #5742 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 22 ++++------ .../backend/tests/test_anthropic_web_fetch.py | 42 +++++++------------ .../src/features/chat/api/chat-adapter.ts | 37 +++++++++------- .../frontend/src/features/chat/chat-page.tsx | 23 ++++++++++ .../features/chat/provider-capabilities.ts | 8 ++-- .../src/features/chat/shared-composer.tsx | 32 +++++++++++++- .../chat/stores/chat-runtime-store.ts | 24 +++++++++++ 7 files changed, 125 insertions(+), 63 deletions(-) diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 0904426633..aff4a31f27 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -1664,25 +1664,19 @@ class ExternalProviderClient: ) body["tools"] = anthropic_tools - # Anthropic server-side web_fetch — see - # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-fetch-tool - # `web_fetch_20250910` reads a single URL (text or PDF) and - # returns a document block in a `web_fetch_tool_result`. For - # safety Anthropic only lets the model fetch URLs that already - # appeared in the conversation (user message, prior tool - # result, web_search hit) — there is no domain restriction we - # have to apply locally. No beta header is required today; the - # tool ships under the standard `2023-06-01` API version. We - # mirror the web_search wiring: max_uses cap, opt in via - # `enabled_tools=["web_fetch"]`, citations off by default - # because the frontend already paints source pills from the - # generic tool_end payload. + # Anthropic server-side web_fetch reads a single URL (text/PDF) + # and returns a `web_fetch_tool_result` document block. Opt in + # via `enabled_tools=["web_fetch"]`; no beta header required. + # `_anthropic_web_fetch_version` picks `web_fetch_20260209` + # (dynamic filtering) for Opus 4.6/4.7 + Sonnet 4.6, falling + # back to `web_fetch_20250910` elsewhere; mismatched variants + # return 400 so the per-model picker is required. web_fetch_enabled = bool(enabled_tools and "web_fetch" in enabled_tools) if web_fetch_enabled: anthropic_tools = list(body.get("tools") or []) anthropic_tools.append( { - "type": "web_fetch_20250910", + "type": _anthropic_web_fetch_version(model), "name": "web_fetch", "max_uses": 5, } diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index 7277757fcf..da10d679eb 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -2,26 +2,12 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Unit tests for Anthropic's server-side `web_fetch_20250910` tool -translation in `_stream_anthropic`. - -Covers: -- Request body: when ``enabled_tools=["web_fetch"]``, the outbound - ``tools`` array carries ``{"type":"web_fetch_20250910", - "name":"web_fetch", "max_uses":5}``. No beta header is required. -- Combined request: ``enabled_tools=["web_search","web_fetch", - "code_execution"]`` sends all three tool entries. -- Disabled by default: with ``enabled_tools=["web_search"]`` (or None), - the body does NOT carry a web_fetch entry. -- SSE translation (success): a `web_fetch` server_tool_use streaming - ``{"url": "..."}`` followed by a `web_fetch_tool_result` block with - a document source emits one ``tool_start`` and one ``tool_end`` - `_toolEvent`. The ``tool_start.arguments.url`` matches the fetched - URL and the ``tool_end.result`` carries the Title / URL / snippet - prefix the source-pill renderer expects. -- SSE translation (error): a `web_fetch_tool_error` with - ``error_code="url_not_accessible"`` renders as ``"Error: - url_not_accessible"`` in the tool_end result. +Unit tests for Anthropic's `web_fetch_20250910` / `web_fetch_20260209` +translation in ``_stream_anthropic``. Covers request body emission +(version picked by ``_anthropic_web_fetch_version``: ``_20260209`` for +Opus 4.6/4.7 + Sonnet 4.6, ``_20250910`` otherwise), combined tool +requests, off-by-default behavior, and SSE translation of success and +``url_not_accessible`` error paths into ``tool_start`` / ``tool_end``. """ import asyncio @@ -117,8 +103,9 @@ def test_web_fetch_tool_appended_to_request_body(monkeypatch): body = captured["body"] tools = body.get("tools") or [] + # claude-opus-4-7 routes web_fetch to _20260209 (dynamic filtering). assert { - "type": "web_fetch_20250910", + "type": "web_fetch_20260209", "name": "web_fetch", "max_uses": 5, } in tools @@ -157,13 +144,10 @@ def test_web_fetch_combined_with_web_search_and_code_execution(monkeypatch): tools = captured["body"].get("tools") or [] tool_types = [t.get("type") for t in tools] - # After PR 5679's per-model tool version dispatch landed, - # claude-opus-4-7 routes web_search to the _20260209 variant and - # code_execution to _20260120. web_fetch still hardcodes - # _20250910 today; see follow-up to thread it through - # _anthropic_web_fetch_version. + # claude-opus-4-7 routes web_search and web_fetch to _20260209 + # and code_execution to _20260120 (per PR 5679 dispatch). assert "web_search_20260209" in tool_types, tool_types - assert "web_fetch_20250910" in tool_types, tool_types + assert "web_fetch_20260209" in tool_types, tool_types assert "code_execution_20260120" in tool_types, tool_types # Code-execution still adds its beta flag; web_fetch must not # have accidentally stripped it. @@ -199,7 +183,9 @@ def test_no_web_fetch_tool_when_pill_off(monkeypatch): _drive(run()) tools = captured["body"].get("tools") or [] - assert all(t.get("type") != "web_fetch_20250910" for t in tools) + assert all( + t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools + ) # ── SSE translation ───────────────────────────────────────────────── diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 800395bddc..d63db25ab7 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -858,7 +858,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Re-read store after potential auto-load / model ready wait runtime = useChatRuntimeStore.getState(); const { params } = runtime; - const { supportsTools, toolsEnabled, codeToolsEnabled, imageToolsEnabled } = runtime; + const { + supportsTools, + toolsEnabled, + codeToolsEnabled, + imageToolsEnabled, + webFetchToolsEnabled, + } = runtime; const externalSelection = parseExternalModelId(params.checkpoint); const isExternalRequest = externalSelection !== null; if ( @@ -914,14 +920,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider.baseUrl, ), ); - // web_fetch shares the Search pill with web_search (no separate - // UI toggle), so it follows toolsEnabled. Anthropic is the only - // provider that ships it today; on others providerSupportsBuiltinWebFetch - // returns false and this stays inert. + // Fetch pill is independent of Search (Anthropic bills web_fetch + // separately from web_search). Sourced from `webFetchToolsEnabled`; + // on providers without web_fetch the toggle is forced off in + // chat-page's runtime setState. const webFetchEnabledForThisTurn = Boolean( externalProvider && - toolsEnabled && + webFetchToolsEnabled && providerSupportsBuiltinWebFetch(externalProvider.providerType), ); const providerShipsWebFetch = Boolean( @@ -983,14 +989,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const webLabel = providerShipsWebFetch ? "web search or web fetch" : "web search"; - if (!webSearchEnabledForThisTurn && !codeExecEnabledForThisTurn) { + // Treat search and fetch as a single "any web tool" axis so + // the guard only warns when neither pill is on; checking + // webSearchEnabledForThisTurn alone mis-fired when only Fetch + // was on and suppressed live web_fetch calls. + const anyWebEnabledForThisTurn = + webSearchEnabledForThisTurn || webFetchEnabledForThisTurn; + if (!anyWebEnabledForThisTurn && !codeExecEnabledForThisTurn) { disabledToolGuard = `You do not have ${webLabel} or code execution tools in this conversation. ` + "Answer from your own knowledge. " + "If a request genuinely requires tool use, live data fetch or running code, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; - } else if (!webSearchEnabledForThisTurn) { + } else if (!anyWebEnabledForThisTurn) { disabledToolGuard = `You do not have ${webLabel} tools in this conversation. ` + "You may still use code execution tools when they are available and useful. " + @@ -1467,13 +1479,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { enable_tools: true, enabled_tools: [ ...(webSearchEnabledForThisTurn ? ["web_search"] : []), - // Pair web_fetch with the Search pill on any - // provider that ships it (Anthropic today). The - // common workflow is "search returns URLs, fetch - // reads them"; without web_fetch the model can - // surface a citation but cannot quote from the - // page body, which is the whole point of the - // tool. There is no separate UI toggle yet. + // web_fetch has its own Fetch pill, independent + // of Search. Anthropic-only today. ...(webFetchEnabledForThisTurn ? ["web_fetch"] : []), ...(codeExecEnabledForThisTurn ? ["code_execution"] : []), // OpenAI Responses-API only: `image_generation` diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 85ed0f7eef..1783a08281 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -57,6 +57,7 @@ import { getProviderCapabilities, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, + providerSupportsBuiltinWebFetch, providerSupportsBuiltinWebSearch, } from "./provider-capabilities"; import { ChatRuntimeProvider } from "./runtime-provider"; @@ -71,6 +72,7 @@ import { CHAT_CODE_TOOLS_ENABLED_KEY, CHAT_IMAGE_TOOLS_ENABLED_KEY, CHAT_TOOLS_ENABLED_KEY, + CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, loadOptionalBool, useChatRuntimeStore, } from "./stores/chat-runtime-store"; @@ -779,6 +781,9 @@ export function ChatPage(): ReactElement { selection.modelId, provider?.baseUrl, ); + const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( + provider?.providerType, + ); // Kimi's k2.6/k2.5 default to thinking enabled on the server side // (per https://platform.kimi.ai/docs/models). Mirror that default // in the UI so the Think pill comes up clicked when the user picks @@ -801,6 +806,9 @@ export function ChatPage(): ReactElement { const storedImageToolsEnabled = loadOptionalBool( CHAT_IMAGE_TOOLS_ENABLED_KEY, ); + const storedWebFetchToolsEnabled = loadOptionalBool( + CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, + ); const nextToolsEnabled = supportsBuiltinWebSearch ? isKimi ? false @@ -834,6 +842,7 @@ export function ChatPage(): ReactElement { supportsBuiltinWebSearch, supportsBuiltinCodeExecution, supportsBuiltinImageGeneration, + supportsBuiltinWebFetch, toolsEnabled: nextToolsEnabled, codeToolsEnabled: supportsBuiltinCodeExecution ? (storedCodeToolsEnabled ?? false) @@ -841,6 +850,10 @@ export function ChatPage(): ReactElement { imageToolsEnabled: supportsBuiltinImageGeneration ? (storedImageToolsEnabled ?? false) : false, + // Default Fetch off (Anthropic bills per fetch); deliberate opt-in. + webFetchToolsEnabled: supportsBuiltinWebFetch + ? (storedWebFetchToolsEnabled ?? false) + : false, }); }, [externalProvidersForChat, inferenceParams.checkpoint]); const canCompare = useMemo(() => { @@ -1008,6 +1021,9 @@ export function ChatPage(): ReactElement { selectedExternal?.modelId, selectedProvider?.baseUrl, ); + const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( + selectedProvider?.providerType, + ); // See sibling useEffect above: Kimi's k2.x default to thinking // enabled, so the Think pill comes up clicked. Search pill stays // off by default; mutual exclusion flips them via the composer. @@ -1026,6 +1042,9 @@ export function ChatPage(): ReactElement { const storedImageToolsEnabled = loadOptionalBool( CHAT_IMAGE_TOOLS_ENABLED_KEY, ); + const storedWebFetchToolsEnabled = loadOptionalBool( + CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, + ); const nextToolsEnabled = supportsBuiltinWebSearch ? isKimi ? false @@ -1067,6 +1086,7 @@ export function ChatPage(): ReactElement { supportsBuiltinWebSearch, supportsBuiltinCodeExecution, supportsBuiltinImageGeneration, + supportsBuiltinWebFetch, toolsEnabled: nextToolsEnabled, codeToolsEnabled: supportsBuiltinCodeExecution ? (storedCodeToolsEnabled ?? false) @@ -1074,6 +1094,9 @@ export function ChatPage(): ReactElement { imageToolsEnabled: supportsBuiltinImageGeneration ? (storedImageToolsEnabled ?? false) : false, + webFetchToolsEnabled: supportsBuiltinWebFetch + ? (storedWebFetchToolsEnabled ?? false) + : false, ...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }), }); return; diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 562a60a18f..1748e098b9 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -123,11 +123,9 @@ export function providerSupportsBuiltinWebSearch( /** * Whether the external provider exposes a server-side web_fetch tool - * that retrieves a single URL (text or PDF) and emits a document block. - * Only Anthropic ships one today (`web_fetch_20250910`); the chat - * composer pairs it with the Search pill because the typical workflow - * is "search returns URLs, fetch reads them" and the UI doesn't (yet) - * expose web_fetch as an independent toggle. + * (single URL, text or PDF) emitting a document block. Anthropic-only + * today (`web_fetch_20250910` / `web_fetch_20260209`). Gates the + * composer's standalone Fetch pill, independent of Search. */ export function providerSupportsBuiltinWebFetch( providerType: string | null | undefined, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 246fd81510..76e77d1288 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -21,7 +21,7 @@ import { isTauri } from "@/lib/api-base"; import { isMultimodalResponse } from "./types/api"; import { getImageInputUnavailableReason } from "./utils/image-input-support"; import { useAui } from "@assistant-ui/react"; -import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; +import { ArrowUpIcon, DownloadIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; import { toast } from "@/lib/toast"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers"; @@ -34,6 +34,7 @@ import { getExternalReasoningCapabilities, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, + providerSupportsBuiltinWebFetch, } from "./provider-capabilities"; import { type CompositionEvent, @@ -336,6 +337,12 @@ export function SharedComposer({ const setImageToolsEnabled = useChatRuntimeStore( (s) => s.setImageToolsEnabled, ); + const webFetchToolsEnabled = useChatRuntimeStore( + (s) => s.webFetchToolsEnabled, + ); + const setWebFetchToolsEnabled = useChatRuntimeStore( + (s) => s.setWebFetchToolsEnabled, + ); const lastOpenRouterChosenModel = useChatRuntimeStore( (s) => s.lastOpenRouterChosenModel, ); @@ -426,6 +433,9 @@ export function SharedComposer({ effectiveExternalModelId, selectedExternalProvider?.baseUrl, ); + const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( + selectedExternalProvider?.providerType, + ); const searchDisabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); const codeDisabled = @@ -437,6 +447,9 @@ export function SharedComposer({ // the pill row stays compact for providers without the capability. const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration; const showImagePill = supportsBuiltinImageGeneration; + // Fetch pill: Anthropic-only (web_fetch_20250910 / web_fetch_20260209). + const webFetchDisabled = !modelLoaded || !supportsBuiltinWebFetch; + const showWebFetchPill = supportsBuiltinWebFetch; // Backwards-compatible alias for any other call site that may still // reference `toolsDisabled` (rare; both pills used it before). const toolsDisabled = codeDisabled; @@ -1106,6 +1119,23 @@ export function SharedComposer({ Images )} + {showWebFetchPill && ( + + )}
{dictationSupported && ( 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 73266b9234..9a6f0c982f 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -25,6 +25,8 @@ export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled"; +export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = + "unsloth_chat_web_fetch_tools_enabled"; // External provider selection is encoded into `params.checkpoint` as // `external::::`. PersistedChatSettings deliberately @@ -262,9 +264,21 @@ type ChatRuntimeStore = { * receive the tool because their runtime cannot dispatch it. */ supportsBuiltinImageGeneration: boolean; + /** + * Whether the active external provider exposes a server-side + * web_fetch tool (Anthropic's `web_fetch_20250910` / + * `web_fetch_20260209`). Gates the composer's Fetch pill, + * independent of Search. + */ + supportsBuiltinWebFetch: boolean; toolsEnabled: boolean; codeToolsEnabled: boolean; imageToolsEnabled: boolean; + /** + * Fetch pill state, independent of `toolsEnabled` (Search). Only + * consulted when `providerSupportsBuiltinWebFetch` is true. + */ + webFetchToolsEnabled: boolean; toolStatus: string | null; generatingStatus: string | null; autoHealToolCalls: boolean; @@ -326,6 +340,7 @@ type ChatRuntimeStore = { setToolsEnabled: (enabled: boolean, options?: { persist?: boolean }) => void; setCodeToolsEnabled: (enabled: boolean) => void; setImageToolsEnabled: (enabled: boolean) => void; + setWebFetchToolsEnabled: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; @@ -567,9 +582,11 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsBuiltinWebSearch: false, supportsBuiltinCodeExecution: false, supportsBuiltinImageGeneration: false, + supportsBuiltinWebFetch: false, toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false), codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false), imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), + webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false), toolStatus: null, generatingStatus: null, autoHealToolCalls: true, @@ -759,9 +776,11 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsBuiltinWebSearch: false, supportsBuiltinCodeExecution: false, supportsBuiltinImageGeneration: false, + supportsBuiltinWebFetch: false, toolsEnabled: false, codeToolsEnabled: false, imageToolsEnabled: false, + webFetchToolsEnabled: false, toolStatus: null, kvCacheDtype: null, loadedKvCacheDtype: null, @@ -821,6 +840,11 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, imageToolsEnabled); return { imageToolsEnabled }; }), + setWebFetchToolsEnabled: (webFetchToolsEnabled) => + set(() => { + saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled); + return { webFetchToolsEnabled }; + }), setToolStatus: (toolStatus) => set({ toolStatus }), setGeneratingStatus: (generatingStatus) => set({ generatingStatus }), setAutoHealToolCalls: (autoHealToolCalls) => From 4854d4579f9d82b995d207c3e8d73115bb77d957 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 23:39:02 -0700 Subject: [PATCH 14/43] Studio: surface Anthropic document citations inline + in Sources panel (#5718) * Studio: surface Anthropic document citations inline + in Sources panel Anthropic's Messages API streams ``citations_delta`` events on ``content_block_delta`` when the request enables ``citations: {enabled: true}`` on document blocks. Each event carries one citation pointing at the source document; previously they were silently dropped, so reader-visible references never reached the chat UI even when the model was citing properly. The proxy now: - dedupes by the type-specific anchor (char_location / page_location / content_block_location / search_result_location) so re-cites of the same span collapse onto a single footnote; - injects ``[N]`` inline right after the matching text run; - forwards the full list as a synthetic ``document_citations`` tool_event at ``message_stop`` so the Sources panel can render per-document footnotes next to web_search / web_fetch citations. Streams that never emit ``citations_delta`` stay byte-identical. References: - https://platform.claude.com/docs/en/build-with-claude/citations - https://platform.claude.com/docs/en/build-with-claude/search-results Tests (5 in test_anthropic_citations.py): passthrough, single char_location, dedup of repeat citations, distinct sources get distinct numbers, search_result_location supported. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: surface Anthropic document_citations in the Sources panel The PR added a backend _toolEvent.type='document_citations' on message_stop and an inline [N] marker in the assistant text, but the chat-adapter only handles container_*/tool_*/sources from web_search and web_fetch tool calls. Reviewers flagged that the inline [N] markers had no matching footnote entries in the Sources panel. Capture the new event into a documentCitationParts buffer, convert each citation dict into a Sources-panel source entry (using document_title or search-result source URL plus cited_text as the snippet), dedupe by id, and append to the final yield alongside the existing web_search/web_fetch sourceParts. * Studio: dedupe search_result_location citations by search_result_index Anthropic's documented search_result_location citation shape carries search_result_index, source, title, and start/end_block_index -- NOT document_index/document_title. The previous key keyed on document_index + document_title + source + start_block_index, so two distinct search results from the same source collapsed onto the same footnote and the second [N] marker was lost. Switch the search_result_location branch to key on the documented fields, and pin the behaviour with a regression test asserting that two citations sharing source/title but with different search_result_index get distinct [1] [2] markers. * Studio: keep each citation distinct across the end-anchor Codex follow-ups on the citations PR: * Backend _anthropic_citation_key now includes the end anchor for every variant (end_char_index, end_page_number, end_block_index). Anthropic ranges are start-AND-end pairs, so a same-start / different-end pair is two distinct citations that previously collapsed onto one footnote. * Frontend documentCitationToSource ids include the position fields (search_result_index, start/end char/page/block) instead of being keyed on URL alone. Two citations from the same document or two search_result_locations with the same source now produce distinct Sources-panel entries, matching the inline [N] numbering. * Studio: key Sources list by per-citation id instead of url Codex flagged that the Sources renderer keys badges on source.url, so two Anthropic document citations sharing the same source URL collide as React keys and one badge gets dropped (or duplicated). The chat-adapter already mints a per-citation id that folds the position fields (search_result_index, start/end char/page/block) into the URL, so the two citations have distinct ids even when their URL matches. Plumb that id through SourceData and use it as the React key for both the measurement badges and the visible SourceBadge list. Falls back to the URL when no id is supplied (web_search and web_fetch source parts). * Studio: enable Anthropic doc citations on input_document blocks Plumb citations: {enabled: true} onto the translated Anthropic document block (both base64 and URL source branches) so the upstream actually emits citations_delta events. Without this opt-in the inline [N] + Sources panel plumbing added in this PR is a no-op for real user PDF / doc uploads. Refs https://platform.claude.com/docs/en/build-with-claude/citations Also add edge-case coverage for the citations_delta path: malformed citations, mixed types per document, reversed indices, missing document_index, non-int block indices, unknown citation type, internal _key never leaking, footnote numbering across content blocks, and the input_document wire-through itself. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reject unsafe citation sources, bound cited_text payload Three follow-ups on top of #5718 surfaced by a deeper review pass: 1) javascript: / data: / vbscript: in citation source is XSS-able. ``documentCitationToSource`` was assigning ``cit.source`` straight into ``Source.url`` and rendering it as an . A hostile model emitting ``cit.source = "javascript:alert(document.domain)"`` would execute on click (openLink only intercepts URLs that contain "://" or start with "mailto:", which both miss the javascript: scheme). Restrict the navigable path to http(s):// only; anything else falls back to the existing #anthropic-doc anchor and the source title still renders the raw identifier for context. Also reject CR/LF inside the URL string. 2) Frontend sources collapse distinct backend footnotes when the citation type differs but positions match. char_location(0,5) and page_location(0,5) over the same source previously deduped into one entry because the id only carried position. Fold citation type into the id anchor so the 1:1 mapping with inline [N] markers is preserved across every citation shape. 3) ``cited_text`` was forwarded unbounded inside the synthetic document_citations tool_event. The Sources panel trims to 240 chars for display anyway; for large RAG / search_result spans (~10kB cited_text is plausible) this inflates SSE bytes 40x for no UI benefit. Truncate server-side at 512 chars with an ellipsis so the description-trim downstream still has room to work and the wire stays bounded. Tests grow from 21 to 22; existing 7 + edge 15 still green. Frontend typecheck clean. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: apply http(s) URL guard to all Sources-panel link sources The previous round only filtered ``cit.source`` inside ``documentCitationToSource``. Two parallel code paths still copied provider/tool-controlled ``URL:`` text directly into clickable ```` Sources-panel links: * ``parseSourcesFromResult`` in chat-adapter.ts (legacy web_search / web_fetch tool result parser) * ``parseSearchResults`` in tool-ui-web-search.tsx (inline tool card) A hostile tool response like ``URL: javascript:alert(1)`` or ``URL: data:text/html,...`` was therefore still rendered as a navigable badge in the Sources panel. Centralise the safe-URL test (``isSafeNavigableSourceUrl``, ``isSafeHttpUrl``) using ``new URL()`` + protocol allowlist + CR/LF rejection, and apply it to both parsers. Unsafe blocks are dropped rather than rewritten to a hash anchor because the web_search / web_fetch parsers have no document-index fallback. Citation conversion now uses the same helper so the in-place http(s) regex and CR/LF check stay in one place. * Shorten citation comments for PR #5718 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/external_provider.py | 119 ++- .../backend/tests/test_anthropic_citations.py | 353 +++++++++ .../tests/test_anthropic_citations_edge.py | 690 ++++++++++++++++++ .../backend/tests/test_multimodal_document.py | 6 + .../src/components/assistant-ui/sources.tsx | 18 +- .../assistant-ui/tool-ui-web-search.tsx | 31 +- .../src/features/chat/api/chat-adapter.ts | 124 +++- 7 files changed, 1326 insertions(+), 15 deletions(-) create mode 100644 studio/backend/tests/test_anthropic_citations.py create mode 100644 studio/backend/tests/test_anthropic_citations_edge.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index aff4a31f27..a1309ae1ac 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -325,6 +325,65 @@ def _anthropic_supports_fast_mode(model: str) -> bool: ) +# Cap on ``cited_text`` forwarded in document_citations tool_events; +# keeps SSE bytes bounded on multi-KB cited spans (frontend trims to +# 240 chars anyway). +_CITED_TEXT_MAX_LEN = 512 + + +def _anthropic_citation_key(citation: dict[str, Any]) -> tuple: + """Stable dedup key for an Anthropic ``citations_delta.citation``. + + Anchor fields vary per type (char_location, page_location, + content_block_location, search_result_location); both start AND + exclusive end indices are part of the key so same-start / + different-end pairs stay distinct. search_result_location keys on + ``search_result_index`` + ``source`` instead of document_index so + distinct results with the same source don't collapse. Unknown + shapes fall back to a stringified copy (more entries, never + collisions). See + https://platform.claude.com/docs/en/build-with-claude/citations + and https://platform.claude.com/docs/en/build-with-claude/search-results. + """ + ctype = citation.get("type") + doc = citation.get("document_index") + title = citation.get("document_title") or "" + if ctype == "char_location": + return ( + ctype, + doc, + title, + citation.get("start_char_index"), + citation.get("end_char_index"), + ) + if ctype == "page_location": + return ( + ctype, + doc, + title, + citation.get("start_page_number"), + citation.get("end_page_number"), + ) + if ctype == "content_block_location": + return ( + ctype, + doc, + title, + citation.get("start_block_index"), + citation.get("end_block_index"), + ) + if ctype == "search_result_location": + return ( + ctype, + citation.get("search_result_index"), + citation.get("source"), + citation.get("title") or "", + citation.get("start_block_index"), + citation.get("end_block_index"), + ) + return (ctype, _json.dumps(citation, sort_keys = True)) + + class _MistralThinkingSpec(NamedTuple): models: tuple[str, ...] style: Literal["prompt_mode", "reasoning_effort", "disabled"] @@ -1460,6 +1519,11 @@ class ExternalProviderClient: "media_type": media_type, "data": b64data, }, + # Opt into Anthropic's natural-citation + # pipeline; without this no citations_delta + # events fire. See + # https://platform.claude.com/docs/en/build-with-claude/citations + "citations": {"enabled": True}, } if title: doc_block["title"] = title @@ -1471,6 +1535,7 @@ class ExternalProviderClient: "type": "url", "url": url, }, + "citations": {"enabled": True}, } if title: doc_block["title"] = title @@ -1925,6 +1990,12 @@ class ExternalProviderClient: # the next turn. current_compaction: Optional[dict[str, Any]] = None compaction_blocks_seen = 0 + # Document citations from ``citations_delta`` events. + # Deduped by type-specific anchor key; inline [N] is + # injected after each cited run, and the full list is + # forwarded as a synthetic document_citations tool_event + # on message_stop for the Sources panel. + document_citations: list[dict[str, Any]] = [] # Counts surfaced in the final log line so reports of # "Code execution did nothing" can be triaged at a # glance. generated_files_count is interesting for the @@ -2276,10 +2347,27 @@ class ExternalProviderClient: thinking_open = False if text: yield _content_chunk(text) - # Citations on text deltas are attached - # per-call by Anthropic via the - # `web_search_tool_result` block; we don't - # need to scrape them off the text events. + # web_search citations: web_search_tool_result. + # User-doc citations: citations_delta below. + elif delta_type == "citations_delta": + # One citation per event; collapse onto a + # numbered footnote list and inject [N] + # inline. See + # https://platform.claude.com/docs/en/build-with-claude/citations + cit = delta.get("citation") + if isinstance(cit, dict): + key = _anthropic_citation_key(cit) + idx_for_marker: Optional[int] = None + for idx, existing in enumerate( + document_citations, start = 1 + ): + if existing.get("_key") == key: + idx_for_marker = idx + break + if idx_for_marker is None: + document_citations.append({**cit, "_key": key}) + idx_for_marker = len(document_citations) + yield _content_chunk(f"[{idx_for_marker}]") elif delta_type == "input_json_delta": # Streamed partial_json carrying tool inputs # — the search query for web_search, or the @@ -2609,6 +2697,29 @@ class ExternalProviderClient: if thinking_open: yield _content_chunk("") thinking_open = False + # Forward document_citations so the Sources + # panel can render the inline [N] footnotes. + # ``cited_text`` is truncated server-side to + # keep SSE bytes bounded on long spans. + if document_citations: + clean_cits = [] + for c in document_citations: + entry = {k: v for k, v in c.items() if k != "_key"} + cited = entry.get("cited_text") + if ( + isinstance(cited, str) + and len(cited) > _CITED_TEXT_MAX_LEN + ): + entry["cited_text"] = ( + cited[:_CITED_TEXT_MAX_LEN] + "…" + ) + clean_cits.append(entry) + yield _emit_tool_event( + { + "type": "document_citations", + "citations": clean_cits, + } + ) # Final include_usage-style chunk so callers can # see cache_creation / cache_read without # scraping the server log. diff --git a/studio/backend/tests/test_anthropic_citations.py b/studio/backend/tests/test_anthropic_citations.py new file mode 100644 index 0000000000..ab5ba10b56 --- /dev/null +++ b/studio/backend/tests/test_anthropic_citations.py @@ -0,0 +1,353 @@ +# 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 Anthropic ``citations_delta`` handling in the streaming proxy. + +Verifies the proxy injects inline ``[N]`` markers after cited text, +dedupes by type-specific anchor (char_location, page_location, +content_block_location, search_result_location), forwards a synthetic +``document_citations`` tool_event at message_stop, and stays inert when +no citations_delta events fire. See +https://platform.claude.com/docs/en/build-with-claude/citations +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _sse(events: list[dict]) -> bytes: + out = [] + for e in events: + ev = e.get("type", "message") + out.append(f"event: {ev}\ndata: {json.dumps(e)}\n\n") + return "".join(out).encode("utf-8") + + +def _capture(monkeypatch, events: list[dict]) -> list[str]: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + lines: list[str] = [] + + async def run(): + client = _make_client() + try: + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "what color is grass?"}], + model = "claude-opus-4-7", + max_tokens = 64, + ): + lines.append(line) + finally: + await client.close() + + _drive(run()) + return lines + + +def _message_start() -> dict: + return { + "type": "message_start", + "message": { + "id": "m1", + "content": [], + "model": "claude-opus-4-7", + "role": "assistant", + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + + +def _content_block_start_text() -> dict: + return { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + } + + +def _text_delta(text: str, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "text_delta", "text": text}, + } + + +def _citations_delta(citation: dict, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "citations_delta", "citation": citation}, + } + + +def _content_block_stop(index: int = 0) -> dict: + return {"type": "content_block_stop", "index": index} + + +def _message_delta_end() -> dict: + return {"type": "message_delta", "delta": {"stop_reason": "end_turn"}} + + +def _message_stop() -> dict: + return {"type": "message_stop"} + + +def _joined(lines: list[str]) -> str: + return "\n".join(lines) + + +def test_no_citations_stream_unchanged(monkeypatch): + """Plain text streams pass through with no inline markers and no + document_citations tool_event.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Grass is green." in body + assert "document_citations" not in body + assert "[1]" not in body + + +def test_single_char_location_emits_inline_marker(monkeypatch): + cit = { + "type": "char_location", + "cited_text": "The grass is green.", + "document_index": 0, + "document_title": "Example", + "start_char_index": 0, + "end_char_index": 20, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Grass is green." in body + assert "[1]" in body, body + assert "document_citations" in body, body + assert '"document_index": 0' in body, body + assert "_key" not in body, body + + +def test_duplicate_citation_dedupes_to_same_number(monkeypatch): + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Example", + "start_char_index": 0, + "end_char_index": 20, + "cited_text": "The grass is green.", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass."), + _citations_delta(cit), + _text_delta(" Still green."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert body.count("[1]") == 2, body + citation_blob = body[body.index("document_citations") :] + assert citation_blob.count('"start_char_index"') == 1, citation_blob + + +def test_distinct_sources_get_distinct_numbers(monkeypatch): + cit1 = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc A", + "start_char_index": 0, + "end_char_index": 5, + } + cit2 = { + "type": "page_location", + "document_index": 1, + "document_title": "Doc B", + "start_page_number": 3, + "end_page_number": 4, + } + cit3 = { + "type": "content_block_location", + "document_index": 2, + "document_title": "Doc C", + "start_block_index": 0, + "end_block_index": 1, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("First"), + _citations_delta(cit1), + _text_delta(" Second"), + _citations_delta(cit2), + _text_delta(" Third"), + _citations_delta(cit3), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body and "[2]" in body and "[3]" in body, body + assert body.index("[1]") < body.index("[2]") < body.index("[3]") + + +def test_search_result_location_supported(monkeypatch): + cit = { + "type": "search_result_location", + "document_index": 0, + "document_title": "Anthropic Search Results", + "source": "https://example.com/doc.html", + "start_block_index": 0, + "end_block_index": 1, + "cited_text": "blah", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Some sourced fact."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body + assert "search_result_location" in body + assert "example.com/doc.html" in body + + +def test_same_start_different_end_offsets_get_distinct_numbers(monkeypatch): + """Same start_char_index + different end_char_index = distinct spans, + so they must get distinct footnote numbers (ranges use exclusive end).""" + cit_a = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 100, + "end_char_index": 150, + "cited_text": "first half", + } + cit_b = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 100, + "end_char_index": 250, + "cited_text": "wider span", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("A "), + _citations_delta(cit_a), + _text_delta(" and B "), + _citations_delta(cit_b), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + assert "[2]" in body, body + + +def test_search_result_location_different_indices_get_distinct_numbers(monkeypatch): + """Same source + different search_result_index = distinct footnotes + (matches the Anthropic search-result citation contract).""" + cit_a = { + "type": "search_result_location", + "search_result_index": 0, + "source": "https://example.com/result.html", + "title": "Result", + "start_block_index": 0, + "end_block_index": 1, + "cited_text": "first", + } + cit_b = { + "type": "search_result_location", + "search_result_index": 1, + "source": "https://example.com/result.html", + "title": "Result", + "start_block_index": 0, + "end_block_index": 1, + "cited_text": "second", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("A "), + _citations_delta(cit_a), + _text_delta(" and B "), + _citations_delta(cit_b), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + assert "[2]" in body, body diff --git a/studio/backend/tests/test_anthropic_citations_edge.py b/studio/backend/tests/test_anthropic_citations_edge.py new file mode 100644 index 0000000000..be1b5f7922 --- /dev/null +++ b/studio/backend/tests/test_anthropic_citations_edge.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case tests for Anthropic ``citations_delta`` handling. + +Complements ``test_anthropic_citations.py``. Covers malformed payloads, +unusual orderings, mixed citation types, and the ``citations: +{enabled: true}`` opt-in attached to translated ``input_document`` +blocks. See +https://platform.claude.com/docs/en/build-with-claude/citations and +https://platform.claude.com/docs/en/build-with-claude/search-results. +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +# ── shared SSE harness ─────────────────────────────────────── + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _sse(events: list[dict]) -> bytes: + out = [] + for e in events: + ev = e.get("type", "message") + out.append(f"event: {ev}\ndata: {json.dumps(e)}\n\n") + return "".join(out).encode("utf-8") + + +def _capture( + monkeypatch, + events: list[dict], + *, + messages: list[dict] | None = None, + captured_body: dict | None = None, +) -> list[str]: + """Drive ``stream_chat_completion`` against a mocked Anthropic + response and return the SSE lines. Pass ``captured_body`` to also + capture the outgoing request body for assertions on the translated + Anthropic shape. + """ + + def handler(request: httpx.Request) -> httpx.Response: + if captured_body is not None: + try: + captured_body.update(json.loads(request.content.decode("utf-8"))) + except Exception: # pragma: no cover -- diagnostic only + pass + return httpx.Response( + 200, + content = _sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + lines: list[str] = [] + + async def run(): + client = _make_client() + try: + async for line in client.stream_chat_completion( + messages = messages + or [{"role": "user", "content": "what color is grass?"}], + model = "claude-opus-4-7", + max_tokens = 64, + ): + lines.append(line) + finally: + await client.close() + + _drive(run()) + return lines + + +def _message_start() -> dict: + return { + "type": "message_start", + "message": { + "id": "m1", + "content": [], + "model": "claude-opus-4-7", + "role": "assistant", + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + + +def _content_block_start_text() -> dict: + return { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + } + + +def _text_delta(text: str, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "text_delta", "text": text}, + } + + +def _citations_delta(citation: dict, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "citations_delta", "citation": citation}, + } + + +def _content_block_stop(index: int = 0) -> dict: + return {"type": "content_block_stop", "index": index} + + +def _message_delta_end() -> dict: + return {"type": "message_delta", "delta": {"stop_reason": "end_turn"}} + + +def _message_stop() -> dict: + return {"type": "message_stop"} + + +def _joined(lines: list[str]) -> str: + return "\n".join(lines) + + +def _citation_payload(body: str) -> dict: + """Pull the ``document_citations`` synthetic tool_event from the + SSE body and return its payload. Raises if absent.""" + assert "document_citations" in body, body + for line in body.splitlines(): + if not line.startswith("data: "): + continue + try: + payload = json.loads(line[len("data: ") :]) + except json.JSONDecodeError: + continue + tool_event = payload.get("_toolEvent") if isinstance(payload, dict) else None + if ( + isinstance(tool_event, dict) + and tool_event.get("type") == "document_citations" + ): + return tool_event + raise AssertionError("document_citations event not parsed out of SSE body") + + +# ── edge cases ─────────────────────────────────────────────── + + +def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch): + """citations_delta before any text_delta must not crash; marker + lands at the start of the block.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "X", + "start_char_index": 0, + "end_char_index": 5, + "cited_text": "x", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _citations_delta(cit), + _text_delta("hello"), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + assert "document_citations" in body, body + + +def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch): + """Non-dict ``delta.citation`` must not crash, emit a marker, or + poison the document_citations list.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Hello."), + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "citations_delta", "citation": "not-a-dict"}, + }, + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Hello." in body + assert "[1]" not in body + assert "document_citations" not in body + + +def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch): + """Missing ``citation`` field is treated like a non-dict citation: + skip without crashing.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Hello."), + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "citations_delta"}, + }, + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Hello." in body + assert "[1]" not in body + assert "document_citations" not in body + + +def test_char_location_with_reversed_indices_does_not_crash(monkeypatch): + """Malformed char_location with reversed indices must not crash; + the dedup key accepts any int pair and still surfaces a footnote.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 300, + "end_char_index": 50, + "cited_text": "?", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Weird."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + payload = _citation_payload(body) + assert payload["citations"][0]["start_char_index"] == 300 + assert payload["citations"][0]["end_char_index"] == 50 + + +def test_page_location_missing_document_index_does_not_crash(monkeypatch): + """page_location missing ``document_index`` still produces a + footnote; dedup key falls back to ``None`` for the missing field.""" + cit = { + "type": "page_location", + "document_title": "Untitled PDF", + "start_page_number": 1, + "end_page_number": 2, + "cited_text": "p1", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("From the PDF:"), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + payload = _citation_payload(body) + assert payload["citations"][0].get("document_index") is None + + +def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypatch): + """content_block_location with string block indices must not crash; + dedup key tolerates non-int values.""" + cit = { + "type": "content_block_location", + "document_index": 0, + "document_title": "Custom", + "start_block_index": "0", + "end_block_index": "1", + "cited_text": "anything", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Cite."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + payload = _citation_payload(body) + assert payload["citations"][0]["start_block_index"] == "0" + + +def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch): + """Unknown citation ``type`` (forward-compat) still dedupes: + identical ones collapse, differing ones get distinct numbers.""" + cit_a = { + "type": "future_shape_location", + "anchor": "abc", + "cited_text": "blah", + } + cit_b = { + "type": "future_shape_location", + "anchor": "xyz", + "cited_text": "blah", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("A"), + _citations_delta(cit_a), + _text_delta(" again"), + _citations_delta(cit_a), + _text_delta(" B"), + _citations_delta(cit_b), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + # cit_a dedupes onto [1], cit_b gets [2]. + assert body.count("[1]") == 2, body + assert body.count("[2]") == 1, body + payload = _citation_payload(body) + assert len(payload["citations"]) == 2 + + +def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch): + """char_location and page_location on the same document_index are + distinct shapes; dedup key uses citation type as its first slot.""" + cit_char = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 10, + } + cit_page = { + "type": "page_location", + "document_index": 0, + "document_title": "Doc", + "start_page_number": 1, + "end_page_number": 2, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("char-cite"), + _citations_delta(cit_char), + _text_delta(" page-cite"), + _citations_delta(cit_page), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body and "[2]" in body, body + payload = _citation_payload(body) + assert len(payload["citations"]) == 2 + + +def test_cited_text_is_preserved_in_synthetic_event(monkeypatch): + """``cited_text`` must survive into the synthetic event so the + Sources panel can render it as a tooltip. Anthropic does not bill + cited_text against output tokens, so preserving it is free.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Trustworthy Doc", + "start_char_index": 0, + "end_char_index": 20, + "cited_text": "The grass is green.", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + payload = _citation_payload(body) + assert payload["citations"][0]["cited_text"] == "The grass is green." + + +def test_internal_key_field_never_leaks_to_client(monkeypatch): + """The internal ``_key`` dedup sentinel must be stripped before + the synthetic event is forwarded; it is not an Anthropic field.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 5, + "cited_text": "..", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("hi"), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + payload = _citation_payload(body) + assert payload["citations"], payload + for c in payload["citations"]: + assert "_key" not in c, c + + +def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch): + """Footnote numbering is per-message, not per-content-block: + citations across separate blocks emit [1] then [2].""" + cit_a = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 5, + } + cit_b = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 100, + "end_char_index": 105, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("first"), + _citations_delta(cit_a, index = 0), + _content_block_stop(0), + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + _text_delta(" second", index = 1), + _citations_delta(cit_b, index = 1), + _content_block_stop(1), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body and "[2]" in body, body + assert body.index("[1]") < body.index("[2]") + payload = _citation_payload(body) + assert len(payload["citations"]) == 2 + + +def test_inline_marker_lands_after_text_run(monkeypatch): + """Inline ``[N]`` must land AFTER the cited text run: Anthropic + streams text then citation, so the proxy emits ``"...green.[1]"`` + not ``"[1]green"``.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 20, + "cited_text": "grass", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _citations_delta(cit), + _text_delta(" Sky is blue."), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + grass = body.index("Grass is green.") + marker = body.index("[1]") + sky = body.index("Sky is blue.") + assert grass < marker < sky, body + + +def test_no_synthetic_event_when_only_text_deltas(monkeypatch): + """No citations_delta means no synthetic ``document_citations`` + event; Sources panel relies on absence to suppress the section.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Just some prose. "), + _text_delta("More prose."), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "document_citations" not in body + assert "[1]" not in body + + +def test_input_document_translation_enables_citations(monkeypatch): + """``input_document`` must translate to an Anthropic ``document`` + block carrying ``citations: {enabled: true}`` (both base64 and url + source branches) so upstream emits citations_delta.""" + captured_b64: dict = {} + _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("ok"), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + messages = [ + { + "role": "user", + "content": [ + { + "type": "input_document", + "file_data": "data:application/pdf;base64,QUJD", + "filename": "spec.pdf", + }, + {"type": "text", "text": "summarise"}, + ], + } + ], + captured_body = captured_b64, + ) + user_msg = captured_b64["messages"][0] + doc_block = next(p for p in user_msg["content"] if p.get("type") == "document") + assert doc_block["source"]["type"] == "base64", doc_block + assert doc_block.get("citations") == {"enabled": True}, doc_block + + captured_url: dict = {} + _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("ok"), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + messages = [ + { + "role": "user", + "content": [ + { + "type": "input_document", + "file_url": "https://example.com/doc.pdf", + "filename": "doc.pdf", + }, + {"type": "text", "text": "summarise"}, + ], + } + ], + captured_body = captured_url, + ) + user_msg = captured_url["messages"][0] + doc_block = next(p for p in user_msg["content"] if p.get("type") == "document") + assert doc_block["source"]["type"] == "url", doc_block + assert doc_block.get("citations") == {"enabled": True}, doc_block + + +# ── cited_text truncation + safe-url citation conversion ──────── + + +def test_cited_text_truncated_in_synthetic_event(monkeypatch): + """``cited_text`` is capped server-side so multi-KB spans do not + balloon the SSE payload.""" + from core.inference.external_provider import _CITED_TEXT_MAX_LEN + + long_quote = "x" * (_CITED_TEXT_MAX_LEN + 4000) + events = [ + { + "type": "message_start", + "message": { + "id": "msg_1", + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "claim "}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "citations_delta", + "citation": { + "type": "char_location", + "document_index": 0, + "document_title": "doc", + "start_char_index": 0, + "end_char_index": 5, + "cited_text": long_quote, + }, + }, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 1}, + }, + {"type": "message_stop"}, + ] + chunks = _capture(monkeypatch, events) + tool_events = [c for c in chunks if "_toolEvent" in c and "document_citations" in c] + assert tool_events, "no document_citations tool event" + payload = json.loads(tool_events[0].split("data: ", 1)[1]) + cited = payload["_toolEvent"]["citations"][0]["cited_text"] + assert len(cited) <= _CITED_TEXT_MAX_LEN + 1, len(cited) + assert cited.endswith("…") diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py index a431b78352..4d7528d238 100644 --- a/studio/backend/tests/test_multimodal_document.py +++ b/studio/backend/tests/test_multimodal_document.py @@ -117,6 +117,8 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch): types = [p.get("type") for p in parts] assert "document" in types, parts doc = _strip_cache(next(p for p in parts if p.get("type") == "document")) + # citations: {enabled: true} opts into Anthropic's natural-citation + # pipeline; without it the citations_delta handler is a no-op. assert doc == { "type": "document", "source": { @@ -124,6 +126,7 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch): "media_type": "application/pdf", "data": _TINY_PDF_B64, }, + "citations": {"enabled": True}, "title": "paper.pdf", } @@ -151,6 +154,7 @@ def test_anthropic_url_pdf_becomes_document_block(monkeypatch): assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, + "citations": {"enabled": True}, } @@ -255,6 +259,7 @@ def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch): assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, + "citations": {"enabled": True}, "title": "doc.pdf", } @@ -283,6 +288,7 @@ def test_anthropic_whitespace_only_data_uri_falls_back_to_file_url(monkeypatch): assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, + "citations": {"enabled": True}, } diff --git a/studio/frontend/src/components/assistant-ui/sources.tsx b/studio/frontend/src/components/assistant-ui/sources.tsx index 3a55c3fa78..140b61f932 100644 --- a/studio/frontend/src/components/assistant-ui/sources.tsx +++ b/studio/frontend/src/components/assistant-ui/sources.tsx @@ -127,6 +127,12 @@ function Source({ // ── Source badge with hover card ───────────────────────────── interface SourceData { + /** + * Stable per-citation key. Two Anthropic document citations into + * different spans of the same source share a ``url``, so React keys + * on ``id`` to keep each footnote distinct. + */ + id: string; url: string; title: string; description?: string; @@ -190,8 +196,14 @@ const SourcesGroup: FC = () => { "url" in part && part.url ) { + const url = part.url as string; + const partId = + typeof (part as { id?: unknown }).id === "string" + ? ((part as { id: string }).id) + : url; sources.push({ - url: part.url as string, + id: partId, + url, title: (part as { title?: string }).title || "", description: (part as { metadata?: { description?: string } }) .metadata?.description, @@ -258,7 +270,7 @@ const SourcesGroup: FC = () => { className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none" > {sources.map((source) => ( - + {source.title || extractDomain(source.url)} @@ -270,7 +282,7 @@ const SourcesGroup: FC = () => { {/* Visible container */}
{displayedSources.map((source) => ( - + ))} {shouldCollapse && !expanded && ( + {isOpen && + typeof document !== "undefined" && + createPortal( + , + document.body, + )} + + ); +} + +function ImageGenerating({ className }: { className?: string }) { + return ( +
+ + Generating image… +
+ ); +} + +function ImageContentFilterError({ + className, + reason, +}: { + className?: string; + reason?: string; +}) { + return ( +
+ +

Image could not be generated

+ {reason &&

{reason}

} +
+ ); +} + +export type ImageActionsProps = { + part: ImageMessagePart; + /** + * Wire to your own generation call to show a regenerate button. The button + * renders only when this is set and the part carries a `prompt`. + */ + onRegenerate?: () => void | Promise; + className?: string; +}; + +function RegenerateButton({ + onRegenerate, +}: { + onRegenerate: () => void | Promise; +}) { + const [isRegenerating, setIsRegenerating] = useState(false); + return ( + + ); +} + +function ImageActions({ part, onRegenerate, className }: ImageActionsProps) { + return ( +
+ + + {onRegenerate && } +
+ ); +} + +const ImageImpl: ImageMessagePartComponent = (props) => { + const { image, filename, status } = props; + const alt = filename || "Image content"; + + if (status?.type === "running") { + return ( + + + {filename} + + ); + } + + if (status?.type === "incomplete" && status.reason === "content-filter") { + return ( + + + + ); + } + + return ( + + + + + {filename} + + ); +}; + +const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & { + Root: typeof ImageRoot; + Preview: typeof ImagePreview; + Filename: typeof ImageFilename; + Zoom: typeof ImageZoom; + Actions: typeof ImageActions; + Generating: typeof ImageGenerating; + ContentFilterError: typeof ImageContentFilterError; +}; + +Image.displayName = "Image"; +Image.Root = ImageRoot; +Image.Preview = ImagePreview; +Image.Filename = ImageFilename; +Image.Zoom = ImageZoom; +Image.Actions = ImageActions; +Image.Generating = ImageGenerating; +Image.ContentFilterError = ImageContentFilterError; + +export { + Image, + ImageRoot, + ImagePreview, + ImageFilename, + ImageZoom, + ImageActions, + ImageGenerating, + ImageContentFilterError, + imageVariants, +}; diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 431e568205..6a99f30508 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -7,6 +7,11 @@ import { UserMessageAttachments, } from "@/components/assistant-ui/attachment"; import { CodeToggleIcon } from "@/components/assistant-ui/code-toggle-icon"; +import { + GeneratedImageOverlayProvider, + useGeneratedImageOverlay, +} from "@/components/assistant-ui/generated-image-overlay-context"; +import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; @@ -39,13 +44,14 @@ import { import { sentAudioNames } from "@/features/chat/api/chat-adapter"; import { parseExternalModelId } from "@/features/chat/external-providers"; import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities"; -import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; +import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; +import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params"; import { isTauri } from "@/lib/api-base"; -import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -79,30 +85,30 @@ import { TerminalIcon, XIcon, } from "lucide-react"; -import { Copy01Icon, Delete02Icon, Edit03Icon, Tick02Icon } from "@hugeicons/core-free-icons"; +import { + Copy01Icon, + Delete02Icon, + Edit03Icon, + Tick02Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ChangeEvent, + type ComponentProps, type CompositionEvent, type FC, - type FormEvent, type KeyboardEvent, useCallback, useEffect, useRef, useState, } from "react"; -import { toast } from "@/lib/toast"; export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean; targetThreadId?: string; -}> = ({ - hideComposer, - hideWelcome, - targetThreadId, -}) => { +}> = ({ hideComposer, hideWelcome, targetThreadId }) => { // Intent-aware autoscroll: replaces assistant-ui's built-in autoscroll // to prevent the streaming-mutation race that makes the viewport snap // back to the bottom while the user is scrolling up (see the hook for @@ -113,85 +119,204 @@ export const Thread: FC<{ const isComposerAttachPending = useAuiState(({ threads }) => targetThreadId ? threads.mainThreadId !== targetThreadId : false, ); + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const threadId = targetThreadId ?? activeThreadId ?? null; return ( - - - - {!hideWelcome && ( - thread.isEmpty && !thread.isLoading}> - - - )} + + + + + {!hideWelcome && ( + thread.isEmpty && !thread.isLoading} + > + + + )} - + - {/* Bottom slack so the last message has breathing room above the + {/* Bottom slack so the last message has breathing room above the sticky scroll-to-bottom button (and the floating composer in single mode). Without this, content would butt against the sticky footer and feel cramped. */} - hideWelcome || !thread.isEmpty}> -
- - - hideWelcome || !thread.isEmpty}> - - - - - - - {!hideComposer && ( - hideWelcome || !thread.isEmpty}> -
+ hideWelcome || !thread.isEmpty}>
-
-
- -
-

- LLMs can make mistakes. Double-check responses. -

-
-
- + + + hideWelcome || !thread.isEmpty}> + + + + + + + + + {!hideComposer && ( + hideWelcome || !thread.isEmpty}> + + + )} + + + + ); +}; + +const GeneratedImageViewportOverlay: FC<{ hideComposer?: boolean }> = ({ + hideComposer, +}) => { + const { overlay, closeOverlay } = useGeneratedImageOverlay(); + + useEffect(() => { + if (!overlay) { + return; + } + document.querySelector(".aui-composer-input")?.focus(); + }, [overlay]); + + if (!overlay) { + return null; + } + + return ( +
+ + +
+
+
+ {overlay.title} +
+
+

+ Generated image +

+ {overlay.metadata ? ( +

+ {overlay.metadata} +

+ ) : null} + {hideComposer ? null : ( +

+ Type edits below, then send. +

+ )} +
+
+ +
+ ); +}; + +const ThreadComposerDock: FC<{ + disabled?: boolean; + threadId?: string | null; +}> = ({ disabled, threadId }) => { + const { overlay } = useGeneratedImageOverlay(); + + return ( +
+
+
+
+ +
+

+ LLMs can make mistakes. Double-check responses. +

+
+
); }; @@ -219,13 +344,17 @@ const ThreadScrollToBottom: FC = () => { ); }; -const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { +const ThreadWelcome: FC<{ + hideComposer?: boolean; + threadId?: string | null; +}> = ({ hideComposer, threadId }) => { const [currentEmoji, setCurrentEmoji] = useState("large sloth drink.png"); useEffect(() => { const hour = new Date().getHours(); if (hour >= 6 && hour < 12) setCurrentEmoji("large sloth drink.png"); - else if (hour >= 12 && hour < 17) setCurrentEmoji("sloth magnify final.png"); + else if (hour >= 12 && hour < 17) + setCurrentEmoji("sloth magnify final.png"); else if (hour >= 17 && hour < 21) setCurrentEmoji("sloth shy large.png"); else setCurrentEmoji("unsloth-gem.png"); }, []); @@ -240,11 +369,7 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
- Sloth mascot + Sloth mascot

Chat with your model

@@ -252,18 +377,21 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { Run GGUFs, safetensors, vision and audio models

- {!hideComposer && } + {!hideComposer && }
); }; -const ComposerAnimated: FC<{ disabled?: boolean }> = ({ disabled }) => { +const ComposerAnimated: FC<{ + disabled?: boolean; + threadId?: string | null; +}> = ({ disabled, threadId }) => { return (
- +
); @@ -293,8 +421,21 @@ const PendingAudioChip: FC = () => { ); }; -const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { - const { inputProps, isComposing, isComposingRef } = useImeComposerInputHandlers(); +const Composer: FC<{ + disabled?: boolean; + threadId?: string | null; +}> = ({ disabled, threadId }) => { + const aui = useAui(); + const { overlay, closeOverlay } = useGeneratedImageOverlay(); + const setImageToolsEnabled = useChatRuntimeStore( + (s) => s.setImageToolsEnabled, + ); + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const setPendingImageEditReference = useChatRuntimeStore( + (s) => s.setPendingImageEditReference, + ); + const { inputProps, isComposing, isComposingRef } = + useImeComposerInputHandlers(); const composerText = useAuiState(({ composer }) => composer.text); const hasAttachments = useAuiState( ({ composer }) => composer.attachments.length > 0, @@ -304,22 +445,78 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { (attachment) => attachment.status.type === "running", ), ); - const hasPendingAudio = useChatRuntimeStore((s) => Boolean(s.pendingAudioName)); + const hasPendingAudio = useChatRuntimeStore((s) => + Boolean(s.pendingAudioName), + ); + const referenceThreadId = threadId ?? activeThreadId ?? null; const hasSendableContent = composerText.trim().length > 0 || hasAttachments || hasPendingAudio; + const shouldBlockSend = useCallback( + () => + !hasSendableContent || isComposingRef.current || hasPendingAttachments, + [hasPendingAttachments, hasSendableContent, isComposingRef], + ); const handleSubmit = useCallback( - (event: FormEvent) => { - if ( - disabled || - !hasSendableContent || - isComposingRef.current || - hasPendingAttachments - ) { + (event: Parameters["onSubmit"]>>[0]) => { + if (disabled || shouldBlockSend()) { event.preventDefault(); + return; + } + + if (overlay) { + const trimmed = composerText.trim(); + if (!trimmed) { + event.preventDefault(); + return; + } + if (!overlay.openaiImageGenerationCallId) { + event.preventDefault(); + toast.error("This generated image cannot be edited", { + description: + "The original image reference is missing. Generate the image again, then retry the edit.", + }); + closeOverlay(); + return; + } + if ((overlay.threadId ?? null) !== referenceThreadId) { + event.preventDefault(); + toast.error("This generated image belongs to another chat", { + description: "Open the original chat and retry the edit.", + }); + closeOverlay(); + return; + } + setImageToolsEnabled(true); + setPendingImageEditReference({ + threadId: overlay.threadId ?? referenceThreadId, + openaiImageGenerationCallId: overlay.openaiImageGenerationCallId, + ...(overlay.openaiResponseId + ? { openaiResponseId: overlay.openaiResponseId } + : {}), + openaiReasoningItem: overlay.openaiReasoningItem, + }); + flushResourcesSync(() => { + aui + .composer() + .setText( + `Use the selected generated image as the reference and apply this edit: ${trimmed}. Preserve everything else exactly.`, + ); + }); + closeOverlay(); } }, - [disabled, hasPendingAttachments, hasSendableContent, isComposingRef], + [ + aui, + closeOverlay, + composerText, + disabled, + overlay, + referenceThreadId, + setImageToolsEnabled, + setPendingImageEditReference, + shouldBlockSend, + ], ); const composerContent = ( @@ -342,11 +539,12 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { /> - !hasSendableContent || isComposingRef.current || hasPendingAttachments + disabled || + !hasSendableContent || + isComposing || + hasPendingAttachments } + shouldBlockSend={shouldBlockSend} /> ); @@ -553,7 +751,6 @@ const ComposerAudioUpload: FC = () => { ); }; - const ReasoningToggle: FC = () => { const modelLoaded = useChatRuntimeStore( (s) => !!s.params.checkpoint && !s.modelLoading, @@ -565,8 +762,12 @@ const ReasoningToggle: FC = () => { const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled); const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle); const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort); - const supportsReasoningOff = useChatRuntimeStore((s) => s.supportsReasoningOff); - const reasoningEffortLevels = useChatRuntimeStore((s) => s.reasoningEffortLevels); + const supportsReasoningOff = useChatRuntimeStore( + (s) => s.supportsReasoningOff, + ); + const reasoningEffortLevels = useChatRuntimeStore( + (s) => s.reasoningEffortLevels, + ); const setReasoningEffort = useChatRuntimeStore((s) => s.setReasoningEffort); const lastOpenRouterChosenModel = useChatRuntimeStore( (s) => s.lastOpenRouterChosenModel, @@ -619,7 +820,8 @@ const ReasoningToggle: FC = () => { effectiveReasoningEnabled && reasoningEffort !== "none"; const disabled = !(modelLoaded && effectiveSupportsReasoning); const formatEffortLabel = (level: typeof reasoningEffort): string => { - if (level !== "xhigh") return level.charAt(0).toUpperCase() + level.slice(1); + if (level !== "xhigh") + return level.charAt(0).toUpperCase() + level.slice(1); const normalized = externalSelection?.modelId?.trim().toLowerCase() ?? ""; if ( normalized.startsWith("claude-opus-4-6") || @@ -677,23 +879,25 @@ const ReasoningToggle: FC = () => { {effectiveReasoningEffortLevels .filter((level) => level !== "none") .map((level) => ( - { - setReasoningEffort(level); - setReasoningEnabled(true); - applyQwenThinkingParams(true); - // Kimi's $web_search builtin forbids thinking, so - // enabling thinking flips the Search pill off. - if (isKimiExternal && toolsEnabled) { - setToolsEnabled(false); - } - }} - > - {formatEffortLabel(level)} - {effectiveReasoningVisualEnabled && reasoningEffort === level ? " \u2713" : ""} - - ))} + { + setReasoningEffort(level); + setReasoningEnabled(true); + applyQwenThinkingParams(true); + // Kimi's $web_search builtin forbids thinking, so + // enabling thinking flips the Search pill off. + if (isKimiExternal && toolsEnabled) { + setToolsEnabled(false); + } + }} + > + {formatEffortLabel(level)} + {effectiveReasoningVisualEnabled && reasoningEffort === level + ? " \u2713" + : ""} + + ))} ); @@ -808,8 +1012,7 @@ const WebSearchToggle: FC = () => { ? externalProviders.find((p) => p.id === externalSelection.providerId) : undefined; const isKimiExternal = selectedExternalProvider?.providerType === "kimi"; - const disabled = - !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); + const disabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); return ( +
+ + +
+
{prompt ? ( -
+
{prompt}
) : null} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index be2799e727..49a5eebd6b 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { getAuthToken } from "@/features/auth/session"; +import { getAuthToken } from "@/features/auth"; import { apiUrl } from "@/lib/api-base"; import { toast } from "@/lib/toast"; import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; @@ -30,12 +30,17 @@ import { providerSupportsBuiltinWebSearch, providerSupportsFastMode, } from "../provider-capabilities"; -import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import { + type PendingImageEditReference, + useChatRuntimeStore, +} from "../stores/chat-runtime-store"; import { useExternalProvidersStore } from "../stores/external-providers-store"; import { isMultimodalResponse } from "../types/api"; import type { OpenAIChatCompletionsRequest, + OpenAIChatMessage, OpenAIMessageContent, + OpenAIReasoningContentPart, } from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; import { getImageInputUnavailableReason } from "../utils/image-input-support"; @@ -399,37 +404,30 @@ function collectImageParts( message: RunMessage, ): Array<{ type: "image_url"; image_url: { url: string } }> { const parts: Array<{ type: "image_url"; image_url: { url: string } }> = []; + const pushImagePart = (part: { type: string }) => { + if (part.type !== "image" || !("image" in part)) { + return; + } + const src = (part as { image: string }).image; + if (!src) { + return; + } + parts.push({ + type: "image_url", + image_url: { + url: src.startsWith("data:") ? src : `data:image/png;base64,${src}`, + }, + }); + }; for (const part of message.content ?? []) { - if (part.type === "image" && "image" in part) { - const src = (part as { image: string }).image; - if (src) { - parts.push({ - type: "image_url", - image_url: { - url: src.startsWith("data:") ? src : `data:image/png;base64,${src}`, - }, - }); - } - } + pushImagePart(part); } if ("attachments" in message && (message.attachments?.length ?? 0) > 0) { for (const attachment of message.attachments ?? []) { for (const part of attachment.content ?? []) { - if (part.type === "image" && "image" in part) { - const src = (part as { image: string }).image; - if (src) { - parts.push({ - type: "image_url", - image_url: { - url: src.startsWith("data:") - ? src - : `data:image/png;base64,${src}`, - }, - }); - } - } + pushImagePart(part); } } } @@ -437,6 +435,66 @@ function collectImageParts( return parts; } +function normalizeOpenAIReasoningItem( + value: unknown, +): OpenAIReasoningContentPart | null { + if (!value || typeof value !== "object") { + return null; + } + const item = value as Record; + if (item.type !== "reasoning" || typeof item.id !== "string" || !item.id) { + return null; + } + const summary = Array.isArray(item.summary) + ? item.summary.flatMap((part) => { + if (!part || typeof part !== "object") { + return []; + } + const summaryPart = part as Record; + return summaryPart.type === "summary_text" && + typeof summaryPart.text === "string" + ? [{ type: "summary_text" as const, text: summaryPart.text }] + : []; + }) + : []; + const normalized: OpenAIReasoningContentPart = { + type: "reasoning", + id: item.id, + summary, + }; + if ( + item.status === "in_progress" || + item.status === "completed" || + item.status === "incomplete" + ) { + normalized.status = item.status; + } + return normalized; +} + +function toOpenAIImageEditReferenceMessage( + reference: PendingImageEditReference, +): OpenAIChatMessage | null { + if (!reference.openaiImageGenerationCallId) { + return null; + } + const content: Exclude = []; + const reasoningItem = normalizeOpenAIReasoningItem( + reference.openaiReasoningItem, + ); + if (reasoningItem) { + content.push(reasoningItem); + } + content.push({ + type: "image_generation_call", + id: reference.openaiImageGenerationCallId, + ...(reference.openaiResponseId + ? { response_id: reference.openaiResponseId } + : {}), + }); + return { role: "assistant", content }; +} + // Refusal flag stamped on assistant metadata when the backend emits the // `anthropic_refusal` _toolEvent. We drop the refused pair from the next // request body (Anthropic guidance: leaving refusals in context keeps @@ -480,10 +538,16 @@ function toOpenAIMessage(message: RunMessage): { if (imageParts.length > 0) { return { role: message.role, - content: [{ type: "text", text: textContent }, ...imageParts], + content: [ + ...(textContent ? [{ type: "text" as const, text: textContent }] : []), + ...imageParts, + ], }; } + if (!textContent) { + return null; + } return { role: message.role, content: textContent }; } @@ -918,17 +982,52 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // the user switches chats while waiting for model load / auto-load. const resolvedThreadId = (unstable_threadId ?? runtime.activeThreadId) || undefined; + const resolvedThreadKey = resolvedThreadId ?? null; + const pendingImageEditReferenceForRun = runtime.pendingImageEditReference; + const selectedImageEditReference = + (pendingImageEditReferenceForRun?.threadId ?? null) === + resolvedThreadKey + ? pendingImageEditReferenceForRun + : null; + const clearSelectedImageEditReference = () => { + if (!selectedImageEditReference) { + return; + } + const store = useChatRuntimeStore.getState(); + const pending = store.pendingImageEditReference; + if ( + pending?.openaiImageGenerationCallId === + selectedImageEditReference.openaiImageGenerationCallId && + pending.openaiResponseId === + selectedImageEditReference.openaiResponseId && + (pending.threadId ?? null) === + (selectedImageEditReference.threadId ?? null) + ) { + store.clearPendingImageEditReference(); + } + }; // Wait for in-progress model load to finish before inferring if (runtime.modelLoading) { toast.info("Waiting for model to finish loading…"); - await waitForModelReady(abortSignal); + try { + await waitForModelReady(abortSignal); + } catch (error) { + clearSelectedImageEditReference(); + throw error; + } } if (!useChatRuntimeStore.getState().params.checkpoint) { // Auto-load the smallest downloaded model - const { loaded, blockedByTrustRemoteCode } = - await autoLoadSmallestModel(); + let loaded: boolean; + let blockedByTrustRemoteCode: boolean; + try { + ({ loaded, blockedByTrustRemoteCode } = await autoLoadSmallestModel()); + } catch (error) { + clearSelectedImageEditReference(); + throw error; + } if (!loaded) { toast.error( blockedByTrustRemoteCode @@ -940,6 +1039,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : "Pick a model in the top bar, then retry.", }, ); + clearSelectedImageEditReference(); throw new Error("Load a model first."); } } @@ -964,6 +1064,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { description: "Turn on Enable connections in Settings → Connections to use hosted models.", }); + clearSelectedImageEditReference(); throw new Error("Connections disabled."); } const externalProvider = isExternalRequest @@ -979,6 +1080,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toast.error("Connection not found.", { description: "Open Settings → Connections and add it again.", }); + clearSelectedImageEditReference(); throw new Error("Connection not found."); } // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. @@ -989,36 +1091,34 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toast.error("Missing API key for selected connection.", { description: "Open Settings → Connections and set the API key again.", }); + clearSelectedImageEditReference(); throw new Error("Missing connection API key."); } - const webSearchEnabledForThisTurn = - Boolean( - externalProvider && - toolsEnabled && - providerSupportsBuiltinWebSearch(externalProvider.providerType), - ); - const codeExecEnabledForThisTurn = - Boolean( - externalProvider && - externalSelection && - codeToolsEnabled && - providerSupportsBuiltinCodeExecution( - externalProvider.providerType, - externalSelection.modelId, - externalProvider.baseUrl, - ), - ); + const webSearchEnabledForThisTurn = Boolean( + externalProvider && + toolsEnabled && + providerSupportsBuiltinWebSearch(externalProvider.providerType), + ); + const codeExecEnabledForThisTurn = Boolean( + externalProvider && + externalSelection && + codeToolsEnabled && + providerSupportsBuiltinCodeExecution( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), + ); // Fetch pill is independent of Search (Anthropic bills web_fetch // separately from web_search). Sourced from `webFetchToolsEnabled`; // on providers without web_fetch the toggle is forced off in // chat-page's runtime setState. - const webFetchEnabledForThisTurn = - Boolean( - externalProvider && - webFetchToolsEnabled && - providerSupportsBuiltinWebFetch(externalProvider.providerType), - ); + const webFetchEnabledForThisTurn = Boolean( + externalProvider && + webFetchToolsEnabled && + providerSupportsBuiltinWebFetch(externalProvider.providerType), + ); const providerShipsWebFetch = Boolean( externalProvider && providerSupportsBuiltinWebFetch(externalProvider.providerType), @@ -1038,6 +1138,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ), ); + if (selectedImageEditReference && !imageGenerationEnabledForThisTurn) { + clearSelectedImageEditReference(); + toast.error("Image editing is unavailable", { + description: + "Select an OpenAI image-generation model, then retry the edit.", + }); + throw new Error("Image generation edit unavailable."); + } + // Two-pass build: a refused assistant turn also drops the user // prompt that triggered it (leaving it in context re-triggers // the classifier). Refusal flag rides assistant @@ -1060,6 +1169,27 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { .filter((message): message is NonNullable => Boolean(message), ); + if (selectedImageEditReference) { + const referenceMessage = toOpenAIImageEditReferenceMessage( + selectedImageEditReference, + ); + if (!referenceMessage) { + clearSelectedImageEditReference(); + toast.error("This generated image cannot be edited", { + description: + "The original image reference is missing. Generate the image again, then retry the edit.", + }); + throw new Error("Generated image edit reference missing."); + } + let insertAt = outboundMessages.length; + for (let i = outboundMessages.length - 1; i >= 0; i -= 1) { + if (outboundMessages[i]?.role === "user") { + insertAt = i; + break; + } + } + outboundMessages.splice(insertAt, 0, referenceMessage); + } const safeSystemPrompt = typeof params.systemPrompt === "string" ? params.systemPrompt : ""; @@ -1084,24 +1214,45 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // was on and suppressed live web_fetch calls. const anyWebEnabledForThisTurn = webSearchEnabledForThisTurn || webFetchEnabledForThisTurn; - if (!anyWebEnabledForThisTurn && !codeExecEnabledForThisTurn) { + if ( + !anyWebEnabledForThisTurn && + !codeExecEnabledForThisTurn && + !imageGenerationEnabledForThisTurn + ) { + disabledToolGuard = + `You do not have ${webLabel}, code execution, or image generation tools in this conversation. ` + + "Answer from your own knowledge. " + + "If a request genuinely requires tool use, live data fetch, running code, or image generation, " + + "inform the user that you do not have access to these capabilities. " + + "Do not return tool-call syntax inside your response."; + } else if (!anyWebEnabledForThisTurn && !codeExecEnabledForThisTurn) { disabledToolGuard = `You do not have ${webLabel} or code execution tools in this conversation. ` + - "Answer from your own knowledge. " + - "If a request genuinely requires tool use, live data fetch or running code, " + + "You may still use image generation tools when they are available and useful. " + + "If a request genuinely requires live data fetch or running code, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; } else if (!anyWebEnabledForThisTurn) { + const availableTools = [ + codeExecEnabledForThisTurn ? "code execution" : null, + imageGenerationEnabledForThisTurn ? "image generation" : null, + ].filter(Boolean); disabledToolGuard = `You do not have ${webLabel} tools in this conversation. ` + - "You may still use code execution tools when they are available and useful. " + + (availableTools.length > 0 + ? `You may still use ${availableTools.join(" and ")} tools when they are available and useful. ` + : "") + "If a request genuinely requires live data fetch or web search tool use, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; } else if (!codeExecEnabledForThisTurn) { + const availableTools = [ + webLabel, + imageGenerationEnabledForThisTurn ? "image generation" : null, + ].filter(Boolean); disabledToolGuard = "You do not have code execution tools in this conversation. " + - `You may still use ${webLabel} tools when they are available and useful. ` + + `You may still use ${availableTools.join(" and ")} tools when they are available and useful. ` + "If a request genuinely requires running code or code execution tool use, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; @@ -1163,6 +1314,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const gatedThreadKey = resolvedThreadId || "__default"; runtime.setThreadRunning(gatedThreadKey, true); runtime.setThreadRunning(gatedThreadKey, false); + clearSelectedImageEditReference(); throw new Error(imageGateReason); } } @@ -1474,8 +1626,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ) { void updateStoredChatThreadEventually(t.id, { openaiCodeExecContainerId: null, - }) - .catch(() => {}); + }).catch(() => {}); continue; } openaiCodeExecContainerId = t.openaiCodeExecContainerId; @@ -1519,8 +1670,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = created.id; void updateStoredChatThreadEventually(resolvedThreadId, { openaiCodeExecContainerId: created.id, - }) - .catch(() => {}); + }).catch(() => {}); } catch { // Fall back to backend's container_auto path on // failure — keeps the chat moving; the next turn @@ -1628,7 +1778,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // attaches `cache_control.ttl` when the value is one of // "5m" / "1h" (see external_provider.py near line 1375), // so unknown values are a no-op end-to-end. - ...(supportsProviderPromptCacheTtl(externalProvider.providerType) && + ...(supportsProviderPromptCacheTtl( + externalProvider.providerType, + ) && (externalProvider.enablePromptCaching ?? true) && isPromptCacheTtl(externalProvider.promptCacheTtl) ? { prompt_cache_ttl: externalProvider.promptCacheTtl } @@ -1706,10 +1858,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let retriedWithRefreshedKey = false; while (true) { try { - const stream = streamChatCompletions( - await buildRequestPayload(retriedWithRefreshedKey), - abortSignal, - ); + let requestPayload: OpenAIChatCompletionsRequest; + try { + requestPayload = await buildRequestPayload(retriedWithRefreshedKey); + } catch (error) { + clearSelectedImageEditReference(); + throw error; + } + clearSelectedImageEditReference(); + const stream = streamChatCompletions(requestPayload, abortSignal); for await (const chunk of stream) { // Handle tool status events @@ -1777,8 +1934,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : "openaiCodeExecContainerId"; void updateStoredChatThreadEventually(resolvedThreadId, { [field]: null, - }) - .catch(() => {}); + }).catch(() => {}); } continue; } @@ -1822,6 +1978,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { size?: string; quality?: string; background?: string; + prompt?: string; }; const imageB64 = toolEvent.image_b64 as string | undefined; if ( @@ -1843,6 +2000,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { size: toolEvent.size as string | undefined, quality: toolEvent.quality as string | undefined, background: toolEvent.background as string | undefined, + prompt: toolEvent.prompt as string | undefined, }; } else if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); @@ -1860,8 +2018,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } else { parsedResult = rawResult; } + const nextArgs = + toolEvent.arguments && + typeof toolEvent.arguments === "object" + ? (toolEvent.arguments as ToolCallMessagePart["args"]) + : undefined; + const mergedArgs = nextArgs + ? { ...(toolCallParts[idx].args ?? {}), ...nextArgs } + : toolCallParts[idx].args; toolCallParts[idx] = { ...toolCallParts[idx], + args: mergedArgs, + argsText: mergedArgs + ? JSON.stringify(mergedArgs) + : toolCallParts[idx].argsText, result: parsedResult, }; } diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 1748e098b9..ef805305be 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -207,15 +207,21 @@ const OPENAI_CODE_EXECUTION_MODEL_PREFIXES = [ /** * Strict check that a provider configuration points at OpenAI's - * managed cloud (api.openai.com), as opposed to a custom OpenAI-compat - * backend (ollama / llama.cpp / vLLM / generic "custom" preset). The - * shell tool ONLY exists on OpenAI cloud; sending it to anything else - * 400s the request. Mirror of the backend's - * `is_openai_cloud = "api.openai.com" in self.base_url` guard. + * managed cloud (api.openai.com) or Azure OpenAI Foundry + * (*.openai.azure.com), as opposed to a custom OpenAI-compat backend + * (ollama / llama.cpp / vLLM / generic "custom" preset). The shell and + * image-generation tools only exist on cloud backends; sending them to + * anything else 400s the request. Mirror of the backend's + * `_is_openai_family_cloud` host check. */ function isOpenAICloudBaseUrl(baseUrl: string | null | undefined): boolean { if (!baseUrl) return true; // No override → uses the default openai.com base. - return baseUrl.trim().toLowerCase().includes("api.openai.com"); + try { + const host = new URL(baseUrl).hostname.toLowerCase(); + return host === "api.openai.com" || host.endsWith(".openai.azure.com"); + } catch { + return false; + } } export function providerSupportsBuiltinCodeExecution( 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 9a6f0c982f..c78f02a474 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -64,6 +64,12 @@ function saveLastExternalCheckpoint(value: string | null): void { } export type ReasoningStyle = "enable_thinking" | "reasoning_effort"; +export type PendingImageEditReference = { + threadId: string | null; + openaiImageGenerationCallId: string; + openaiResponseId?: string; + openaiReasoningItem?: unknown; +}; export type ReasoningEffort = | "none" | "minimal" @@ -300,6 +306,7 @@ type ChatRuntimeStore = { settingsPanelOpen: boolean; pendingAudioBase64: string | null; pendingAudioName: string | null; + pendingImageEditReference: PendingImageEditReference | null; contextUsage: { promptTokens: number; completionTokens: number; @@ -353,6 +360,10 @@ type ChatRuntimeStore = { setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; + setPendingImageEditReference: ( + reference: PendingImageEditReference | null, + ) => void; + clearPendingImageEditReference: () => void; setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void; }; @@ -607,6 +618,7 @@ export const useChatRuntimeStore = create((set, get) => ({ settingsPanelOpen: false, pendingAudioBase64: null, pendingAudioName: null, + pendingImageEditReference: null, contextUsage: null, modelLoading: false, activeNativePathToken: null, @@ -793,6 +805,7 @@ export const useChatRuntimeStore = create((set, get) => ({ defaultChatTemplate: null, chatTemplateOverride: null, loadedChatTemplateOverride: null, + pendingImageEditReference: null, })); }, setReasoningEnabled: (reasoningEnabled, options) => @@ -884,5 +897,9 @@ export const useChatRuntimeStore = create((set, get) => ({ set({ pendingAudioBase64: base64, pendingAudioName: name }), clearPendingAudio: () => set({ pendingAudioBase64: null, pendingAudioName: null }), + setPendingImageEditReference: (pendingImageEditReference) => + set({ pendingImageEditReference }), + clearPendingImageEditReference: () => + set({ pendingImageEditReference: null }), setContextUsage: (contextUsage) => set({ contextUsage }), })); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index f18407413d..b7a61d24b6 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -192,12 +192,31 @@ export interface AudioGenerationResponse { }>; } -export type OpenAIMessageContent = - | string - | Array< - | { type: "text"; text: string } - | { type: "image_url"; image_url: { url: string } } - >; +export type OpenAIReasoningSummaryPart = { + type: "summary_text"; + text: string; +}; + +export type OpenAIReasoningContentPart = { + type: "reasoning"; + id: string; + summary: OpenAIReasoningSummaryPart[]; + status?: "in_progress" | "completed" | "incomplete"; +}; + +export type OpenAIImageGenerationCallContentPart = { + type: "image_generation_call"; + id: string; + response_id?: string; +}; + +export type OpenAIMessageContentPart = + | { type: "text"; text: string } + | { type: "image_url"; image_url: { url: string } } + | OpenAIReasoningContentPart + | OpenAIImageGenerationCallContentPart; + +export type OpenAIMessageContent = string | OpenAIMessageContentPart[]; export interface OpenAIChatMessage { role: "system" | "user" | "assistant"; diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 8f132cd95b..8a1fe13678 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1188,6 +1188,53 @@ border-color: var(--border) !important; } +.generated-image-loading-card { + position: relative; + overflow: hidden; + contain: paint; +} + +.generated-image-loading-wave { + position: relative; + display: grid; + grid-template-columns: repeat(8, minmax(0, 1fr)); + gap: 14px; + width: min(66%, 18rem); + padding: 1.5rem; + border-radius: 1.5rem; +} + +.generated-image-loading-dot { + width: 7px; + height: 7px; + border-radius: 9999px; + background: color-mix(in oklch, var(--muted-foreground) 82%, var(--primary)); + opacity: 0.12; + transform: translate3d(0, 4px, 0) scale(0.72); + animation: generated-image-dot-wave 1850ms var(--ease-out-quart) infinite; + animation-delay: calc((var(--dot-row) * 72ms) + (var(--dot-col) * 72ms)); + will-change: transform, opacity; +} + +@keyframes generated-image-dot-wave { + 0%, + 22%, + 100% { + opacity: 0.1; + transform: translate3d(0, 4px, 0) scale(0.72); + } + + 46% { + opacity: 0.46; + transform: translate3d(0, -3px, 0) scale(0.96); + } + + 66% { + opacity: 0.2; + transform: translate3d(0, 0, 0) scale(0.82); + } +} + /* * prefers-reduced-motion: honour the OS-level "reduce motion" preference. * Tailwind animate-in/out, Radix open/close transforms, infinite shine/pulse From 31ac558a73fd76bc200a4a22f10ceae5a73b9ae4 Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Tue, 26 May 2026 11:37:24 +0200 Subject: [PATCH 17/43] Recipe Studio local model selector (#5769) * feat(recipes): round-trip local model variants * feat(recipes): add local model selector * feat(recipes): wire selector into model editors * fix(recipes): clear stale model state on relink * feat(recipes): load selected local models for jobs * chore(frontend): simplify biome scripts * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(recipes): handle local selector edge cases * fix(recipes): polish local model selector behavior * fix(recipes): delay local model restore until terminal runs * fix(recipes): accept resolved default gguf variants --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/data_recipe/service.py | 41 +- studio/backend/models/inference.py | 6 +- studio/backend/routes/data_recipe/jobs.py | 121 +++- studio/backend/routes/inference.py | 16 +- studio/frontend/package.json | 4 +- studio/frontend/src/features/chat/index.ts | 8 + .../features/chat/presets/preset-policy.ts | 2 +- .../frontend/src/features/chat/types/api.ts | 3 +- .../components/inline/inline-model.tsx | 80 ++- .../models/local-recipe-model-selector.tsx | 644 ++++++++++++++++++ .../dialogs/models/model-config-dialog.tsx | 86 ++- .../easy/github-crawler-easy-view.tsx | 85 ++- .../recipe-studio/executions/tracker.ts | 101 ++- .../hooks/use-recipe-executions.ts | 565 ++++++++++++--- .../stores/helpers/reference-sync.ts | 20 +- .../recipe-studio/stores/recipe-studio.ts | 98 +-- .../src/features/recipe-studio/types/index.ts | 2 + .../utils/graph/recipe-graph-connection.ts | 65 +- .../utils/import/parsers/model-parser.ts | 13 +- .../utils/payload/builders-model.ts | 73 +- .../recipe-studio/utils/payload/validate.ts | 31 +- 21 files changed, 1714 insertions(+), 350 deletions(-) create mode 100644 studio/frontend/src/features/recipe-studio/dialogs/models/local-recipe-model-selector.tsx diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 7fff36aefd..b4ec0ccd94 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -109,7 +109,7 @@ def _apply_data_designer_image_context_patch() -> None: return try: - from data_designer.config.models import ImageContext + from data_designer.config.models import ImageContext # pyright: ignore[reportMissingImports] except ImportError: return @@ -131,7 +131,7 @@ def _apply_data_designer_image_context_patch() -> None: def build_model_providers(recipe: dict[str, Any]): - from data_designer.config.models import ModelProvider + from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] providers: list[ModelProvider] = [] for provider in recipe.get("model_providers", []): @@ -174,7 +174,7 @@ def _validate_recipe_runtime_support( def build_mcp_providers( recipe: dict[str, Any], ) -> list: - from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider + from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] providers: list[MCPProvider | LocalStdioMCPProvider] = [] for provider in recipe.get("mcp_providers", []): @@ -214,16 +214,42 @@ def build_mcp_providers( return providers +def _strip_frontend_model_config_metadata(recipe: dict[str, Any]) -> dict[str, Any]: + model_configs = recipe.get("model_configs") + if not isinstance(model_configs, list): + return recipe + + changed = False + next_model_configs: list[Any] = [] + for model_config in model_configs: + if isinstance(model_config, dict) and "gguf_variant" in model_config: + next_model_config = dict(model_config) + next_model_config.pop("gguf_variant", None) + next_model_configs.append(next_model_config) + changed = True + continue + next_model_configs.append(model_config) + + if not changed: + return recipe + + return { + **recipe, + "model_configs": next_model_configs, + } + + def build_config_builder(recipe: dict[str, Any]): _apply_data_designer_image_context_patch() - from data_designer.config import DataDesignerConfigBuilder - from data_designer.config.processors import ProcessorType + from data_designer.config import DataDesignerConfigBuilder # pyright: ignore[reportMissingImports] + from data_designer.config.processors import ProcessorType # pyright: ignore[reportMissingImports] recipe_core = { key: value for key, value in recipe.items() if key not in {"model_providers", "mcp_providers"} } + recipe_core = _strip_frontend_model_config_metadata(recipe_core) recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators( recipe_core ) @@ -256,8 +282,9 @@ def create_data_designer( artifact_path: str | None = None, ): _apply_data_designer_image_context_patch() - from data_designer.interface.data_designer import DataDesigner + from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports] + recipe = _strip_frontend_model_config_metadata(recipe) model_providers = build_model_providers(recipe) _validate_recipe_runtime_support(recipe, model_providers) @@ -265,7 +292,7 @@ def create_data_designer( # when the pipeline contains no LLM columns. Supply a lightweight stub # so sampler/expression-only recipes can run without a real provider. if not model_providers: - from data_designer.config.models import ModelProvider + from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] model_providers = [ ModelProvider( diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index de2c166d91..0af9425fdc 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -299,7 +299,11 @@ class InferenceStatusResponse(BaseModel): """Current inference backend status""" active_model: Optional[str] = Field( - None, description = "Currently active model identifier" + None, description = "Currently active model display identifier" + ) + model_identifier: Optional[str] = Field( + None, + description = "Loadable identifier for the active model.", ) is_vision: bool = Field( False, description = "Whether the active model is a vision model" diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index da6416e324..107a1657f3 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -95,6 +95,89 @@ def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]: return aliases +def _used_local_model_selections( + recipe: dict[str, Any], local_provider_names: set[str] +) -> dict[tuple[str, str], list[str]]: + used_aliases = _used_llm_model_aliases(recipe) + selections: dict[tuple[str, str], list[str]] = {} + for mc in recipe.get("model_configs", []): + if not isinstance(mc, dict): + continue + alias = mc.get("alias") + if not isinstance(alias, str) or alias not in used_aliases: + continue + provider = mc.get("provider") + if not isinstance(provider, str) or provider not in local_provider_names: + continue + model = mc.get("model") + target = model.strip() if isinstance(model, str) else "" + if not target or target.lower() == "local": + continue + variant = mc.get("gguf_variant") + gguf_variant = variant.strip() if isinstance(variant, str) else "" + selections.setdefault((target, gguf_variant), []).append(alias) + return selections + + +def _single_used_local_model_selection( + recipe: dict[str, Any], local_provider_names: set[str] +) -> tuple[str, str] | None: + selections = _used_local_model_selections(recipe, local_provider_names) + if not selections: + return None + if len(selections) > 1: + aliases = ", ".join(alias for values in selections.values() for alias in values) + raise ValueError( + "Recipes supports one active local model per run. " + f"Select the same local model and GGUF variant for: {aliases}." + ) + return next(iter(selections)) + + +def _loaded_local_model_identity() -> tuple[bool, str, str]: + from routes.inference import get_llama_cpp_backend + from core.inference import get_inference_backend + + llama = get_llama_cpp_backend() + if llama.is_loaded: + model = str(getattr(llama, "model_identifier", "") or "").strip() + variant = str(getattr(llama, "hf_variant", "") or "").strip() + return True, model, variant + + backend = get_inference_backend() + active_model = str(getattr(backend, "active_model_name", "") or "").strip() + if active_model: + return True, active_model, "" + return False, "", "" + + +def _ensure_selected_local_model_loaded( + recipe: dict[str, Any], local_provider_names: set[str] +) -> None: + model_loaded, active_model, active_variant = _loaded_local_model_identity() + if not model_loaded: + raise ValueError( + "No model loaded in Chat. Load a model first, then run the recipe." + ) + + selection = _single_used_local_model_selection(recipe, local_provider_names) + if selection is None: + return + + target, gguf_variant = selection + variant_matches = not gguf_variant or active_variant == gguf_variant + if active_model.lower() != target.lower() or not variant_matches: + selected = f"{target} ({gguf_variant})" if gguf_variant else target + active = ( + f"{active_model} ({active_variant})" if active_variant else active_model + ) + raise ValueError( + "Selected local model is not loaded. " + f"Selected {selected}; active {active or 'none'}. " + "Load the selected model again, then run the recipe." + ) + + def _inject_local_structured_response_format( recipe: dict[str, Any], local_provider_names: set[str] ) -> None: @@ -238,24 +321,12 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona token = "" internal_key_id: Optional[int] = None if local_names & referenced_providers: - # Verify a model is loaded. - # NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded - # or swapped after this check but before the recipe subprocess calls /v1. - # The inference endpoint returns a clear 400 in that case. - # - # Imports are deferred to avoid circular dependencies with inference modules. - from routes.inference import get_llama_cpp_backend - from core.inference import get_inference_backend - - llama = get_llama_cpp_backend() - model_loaded = llama.is_loaded - if not model_loaded: - backend = get_inference_backend() - model_loaded = bool(backend.active_model_name) - if not model_loaded: - raise ValueError( - "No model loaded in Chat. Load a model first, then run the recipe." - ) + # Verify the selected local model is loaded before minting a workflow + # key. This still remains a point-in-time singleton-backend check + # (TOCTOU): a future generation token should bind frontend load and + # job creation, and the inference endpoint returns a clear 400 if the + # model is later unloaded or swapped before the subprocess calls /v1. + _ensure_selected_local_model_loaded(recipe, local_names) from auth import storage # deferred: avoids circular import @@ -287,12 +358,12 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona providers[i].pop("extra_body", None) # Force skip_health_check on any model_config that references a local - # provider. The local /v1/models endpoint only lists the real loaded - # model (e.g. "unsloth/llama-3.2-1b") and not the placeholder "local" - # that the recipe sends as the model id, so data_designer's pre-flight - # health check would otherwise fail before the first completion call. - # The backend route ignores the model id field in chat completions, so - # skipping the check is safe. + # provider. The frontend now sends the explicit selected local model id, + # but llama-server's /v1/models response can still differ from that id + # for local paths, cache aliases, and GGUF variant loads. The recipe run + # has already gated on a loaded local inference backend above, so the + # data_designer model-list health check would be redundant and can reject + # valid local selections. for mc in recipe.get("model_configs", []): if not isinstance(mc, dict): continue @@ -319,7 +390,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona tpl_kwargs = extra_body.get("chat_template_kwargs") if not isinstance(tpl_kwargs, dict): tpl_kwargs = {} - tpl_kwargs.setdefault("enable_thinking", False) + tpl_kwargs["enable_thinking"] = False extra_body["chat_template_kwargs"] = tpl_kwargs params["extra_body"] = extra_body diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9621d18801..a156f2397c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -606,11 +606,16 @@ async def load_model( backend = get_inference_backend() llama_backend = get_llama_cpp_backend() - if request.gguf_variant: + is_direct_gguf_request = model_identifier.lower().endswith(".gguf") + if request.gguf_variant or is_direct_gguf_request: + gguf_variant_matches = is_direct_gguf_request or bool( + llama_backend.hf_variant + and request.gguf_variant + and llama_backend.hf_variant.lower() == request.gguf_variant.lower() + ) if ( llama_backend.is_loaded - and llama_backend.hf_variant - and llama_backend.hf_variant.lower() == request.gguf_variant.lower() + and gguf_variant_matches and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() # Match runtime settings too so Apply isn't dropped (#5401). @@ -619,7 +624,8 @@ async def load_model( and getattr(llama_backend, "_audio_probed", True) ): logger.info( - f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload" + "Model already loaded (GGUF): " + f"{model_log_label} variant={request.gguf_variant or llama_backend.hf_variant}, skipping reload" ) inference_config = load_inference_config(llama_backend.model_identifier) @@ -1373,6 +1379,7 @@ async def get_status( _audio_type = getattr(llama_backend, "_audio_type", None) return InferenceStatusResponse( active_model = _display_model_id, + model_identifier = None if _native_grant_backed else _model_id, is_vision = llama_backend.is_vision, is_gguf = True, gguf_variant = llama_backend.hf_variant, @@ -1435,6 +1442,7 @@ async def get_status( return InferenceStatusResponse( active_model = backend.active_model_name, + model_identifier = backend.active_model_name, is_vision = is_vision, is_gguf = False, is_audio = is_audio, diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 061a2b517d..83b1fd96f9 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -12,8 +12,8 @@ "lint": "eslint .", "preview": "vite preview", "typecheck": "tsc -b --pretty false", - "biome:check": "biome check .", - "biome:fix": "biome check . --write" + "biome:check": "biome check", + "biome:fix": "biome check --write" }, "dependencies": { "@assistant-ui/core": "0.1.17", diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 4726b11fcf..883dea3f3a 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -2,6 +2,14 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export { ChatPage } from "./chat-page"; +export { + getInferenceStatus, + listGgufVariants, + listLocalModels, + loadModel, + type LocalModelInfo, +} from "./api/chat-api"; +export type { GgufVariantDetail } from "./types/api"; export { ChatSettingsPanel, defaultInferenceParams, diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts index ae8c1f41ae..4efbb74f11 100644 --- a/studio/frontend/src/features/chat/presets/preset-policy.ts +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -248,7 +248,7 @@ interface BackendInferenceDefaults { export interface BackendInferenceEnvelope { is_gguf?: boolean; context_length?: number | null; - inference?: BackendInferenceDefaults; + inference?: BackendInferenceDefaults | null; } export function mergeBackendRecommendedInference({ diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index b7a61d24b6..5238875b71 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -143,6 +143,7 @@ export interface UnloadModelRequest { export interface InferenceStatusResponse { active_model: string | null; + model_identifier?: string | null; is_vision: boolean; is_gguf?: boolean; gguf_variant?: string | null; @@ -158,7 +159,7 @@ export interface InferenceStatusResponse { min_p?: number; presence_penalty?: number; trust_remote_code?: boolean; - }; + } | null; requires_trust_remote_code?: boolean; supports_reasoning?: boolean; reasoning_style?: "enable_thinking" | "reasoning_effort"; diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx index 16e99f4fae..1d98d18b56 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx @@ -3,6 +3,7 @@ import { Input } from "@/components/ui/input"; import type { ReactElement } from "react"; +import { LocalRecipeModelSelector } from "../../dialogs/models/local-recipe-model-selector"; import type { ModelConfig, ModelProviderConfig } from "../../types"; import { InlineField } from "./inline-field"; @@ -32,7 +33,9 @@ export function InlineModel(props: InlineModelProps): ReactElement { className="nodrag h-8 w-full text-xs" placeholder="https://api.example.com/v1" value={props.config.endpoint} - onChange={(event) => props.onUpdate({ endpoint: event.target.value })} + onChange={(event) => + props.onUpdate({ endpoint: event.target.value }) + } /> @@ -53,23 +56,32 @@ export function InlineModel(props: InlineModelProps): ReactElement { } // model_config branch - mirror the local-aware provider sync from the - // dialog path so inline edits do not leave stale "local" placeholders - // on external providers and fill the placeholder when switching to local. + // dialog path so inline edits clear stale local-only metadata without + // synthesizing the legacy "local" placeholder. const localNames = props.localProviderNames ?? new Set(); const modelConfig = props.config; - const handleProviderChange = (nextProvider: string) => { - const isLocal = localNames.has(nextProvider); - if (isLocal && !modelConfig.model.trim()) { - props.onUpdate({ provider: nextProvider, model: "local" }); - return; - } - if (!isLocal && modelConfig.model === "local") { - props.onUpdate({ provider: nextProvider, model: "" }); - return; - } - props.onUpdate({ provider: nextProvider }); - }; const isLinkedToLocal = localNames.has(modelConfig.provider); + const handleProviderChange = (nextProvider: string) => { + const nextIsLocal = localNames.has(nextProvider); + if (isLinkedToLocal !== nextIsLocal) { + props.onUpdate({ + provider: nextProvider, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }); + return; + } + props.onUpdate({ + provider: nextProvider, + ...(nextIsLocal + ? {} + : { + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }), + }); + }; return (
@@ -82,12 +94,38 @@ export function InlineModel(props: InlineModelProps): ReactElement { /> - props.onUpdate({ model: event.target.value })} - /> + {isLinkedToLocal ? ( + + props.onUpdate({ + model, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: variant ?? undefined, + }) + } + /> + ) : ( + + props.onUpdate({ + model: event.target.value, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }) + } + /> + )} void; + inputId?: string; + disabled?: boolean; + compact?: boolean; + className?: string; +}; + +function normalizeForSearch(value: string): string { + return value.toLowerCase().replace(/[\s_.-]/g, ""); +} + +function hasGgufSuffix(value: string | null | undefined): boolean { + return GGUF_SUFFIX_PATTERN.test(value ?? ""); +} + +function getModelLabel(model: LocalModelInfo): string { + return model.model_id?.trim() || model.display_name || model.id; +} + +function isDirectGguf(model: LocalModelInfo): boolean { + return model.path.toLowerCase().endsWith(".gguf"); +} + +function isExpandableGguf(model: LocalModelInfo): boolean { + return ( + !isDirectGguf(model) && + (hasGgufSuffix(model.id) || + hasGgufSuffix(model.display_name) || + hasGgufSuffix(model.model_id)) + ); +} + +function sourceLabel(model: LocalModelInfo): string { + switch (model.source) { + case "models_dir": + return "Models"; + case "hf_cache": + return "HF cache"; + case "lmstudio": + return "LM Studio"; + case "custom": + return "Custom folder"; + default: + return "Local"; + } +} + +type SelectedModelSummary = { + label: string; + source: string; + isGguf: boolean; +}; + +function getSelectedModelSummary( + value: string, + selectedModel: LocalModelInfo | null, + ggufVariant?: string | null, +): SelectedModelSummary { + if (!selectedModel) { + return { + label: value, + source: "Local model", + isGguf: Boolean(ggufVariant), + }; + } + + return { + label: getModelLabel(selectedModel), + source: sourceLabel(selectedModel), + isGguf: isDirectGguf(selectedModel) || isExpandableGguf(selectedModel), + }; +} + +function LocalGgufVariantList({ + repoId, + selectedVariant, + onSelect, +}: { + repoId: string; + selectedVariant?: string | null; + onSelect: (variant: string) => void; +}): ReactElement { + const [variants, setVariants] = useState(null); + const [defaultVariant, setDefaultVariant] = useState(null); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + + useEffect(() => { + let cancelled = false; + listGgufVariants(repoId) + .then((response) => { + if (cancelled) { + return; + } + setVariants(response.variants); + setDefaultVariant(response.default_variant); + }) + .catch((err) => { + if (cancelled) { + return; + } + setError( + err instanceof Error ? err.message : "Failed to load variants.", + ); + }) + .finally(() => { + if (!cancelled) { + setLoading(false); + } + }); + return () => { + cancelled = true; + }; + }, [repoId]); + + const sortedVariants = useMemo(() => { + if (!variants) { + return null; + } + return [...variants].sort((a, b) => { + if (a.quant === defaultVariant) { + return -1; + } + if (b.quant === defaultVariant) { + return 1; + } + if (a.downloaded !== b.downloaded) { + return a.downloaded ? -1 : 1; + } + return a.quant.localeCompare(b.quant); + }); + }, [defaultVariant, variants]); + + if (loading) { + return ( +
+ + Loading quantizations... +
+ ); + } + + if (error) { + return
{error}
; + } + + if (!sortedVariants || sortedVariants.length === 0) { + return ( +
+ No GGUF quantizations found for this model. +
+ ); + } + + return ( +
+
+ Quantization +
+
+ {sortedVariants.map((variant) => { + const selected = selectedVariant === variant.quant; + return ( + + ); + })} +
+
+ ); +} + +type SelectorTriggerProps = ComponentPropsWithoutRef<"button"> & { + value: string; + selectedModel: LocalModelInfo | null; + ggufVariant?: string | null; + inputId?: string; + disabled: boolean; + compact: boolean; + className?: string; +}; + +const SelectorTrigger = forwardRef( + function SelectorTrigger( + { + value, + selectedModel, + ggufVariant, + inputId, + disabled, + compact, + className, + ...triggerProps + }, + ref, + ): ReactElement { + const selected = getSelectedModelSummary(value, selectedModel, ggufVariant); + + return ( + + ); + }, +); + +function LocalModelRow({ + model, + selected, + expanded, + probing, + ggufVariant, + onSelectModel, + onSelectVariant, +}: { + model: LocalModelInfo; + selected: boolean; + expanded: boolean; + probing: boolean; + ggufVariant?: string | null; + onSelectModel: (model: LocalModelInfo) => void; + onSelectVariant: (modelId: string, variant: string) => void; +}): ReactElement { + const expandable = isExpandableGguf(model); + const directGguf = isDirectGguf(model); + + return ( +
+ + {expanded ? ( + onSelectVariant(model.id, variant)} + /> + ) : null} +
+ ); +} + +function LocalModelResults({ + loading, + error, + models, + value, + ggufVariant, + expandedModelId, + probingVariantModelId, + onRefresh, + onSelectModel, + onSelectVariant, +}: { + loading: boolean; + error: string | null; + models: LocalModelInfo[]; + value: string; + ggufVariant?: string | null; + expandedModelId: string | null; + probingVariantModelId: string | null; + onRefresh: () => void; + onSelectModel: (model: LocalModelInfo) => void; + onSelectVariant: (modelId: string, variant: string) => void; +}): ReactElement { + if (loading) { + return ( +
+ + Scanning local models... +
+ ); + } + + if (error) { + return ( +
+

{error}

+ +
+ ); + } + + if (models.length === 0) { + return ( +
+

No local models found.

+

+ Download a model or add a scan folder from Chat, then refresh this + list. +

+ + Open Chat model picker + +
+ ); + } + + return ( +
+ {models.map((model) => ( + + ))} +
+ ); +} + +export function LocalRecipeModelSelector({ + value, + ggufVariant, + onChange, + inputId, + disabled = false, + compact = false, + className, +}: LocalRecipeModelSelectorProps): ReactElement { + const [open, setOpen] = useState(false); + const [query, setQuery] = useState(""); + const [models, setModels] = useState([]); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [expandedModelId, setExpandedModelId] = useState(null); + const [probingVariantModelId, setProbingVariantModelId] = useState< + string | null + >(null); + const [refreshKey, setRefreshKey] = useState(0); + + const requestModelRefresh = useCallback(() => { + setLoading(true); + setError(null); + setRefreshKey((key) => key + 1); + }, []); + + const handleOpenChange = useCallback( + (nextOpen: boolean) => { + setOpen(nextOpen); + if (nextOpen) { + requestModelRefresh(); + } + }, + [requestModelRefresh], + ); + + useEffect(() => { + if (!open || refreshKey < 0) { + return; + } + let cancelled = false; + listLocalModels() + .then((response) => { + if (cancelled) { + return; + } + setModels(response.models); + }) + .catch((err) => { + if (cancelled) { + return; + } + setError( + err instanceof Error ? err.message : "Failed to list local models.", + ); + }) + .finally(() => { + if (!cancelled) { + setLoading(false); + } + }); + return () => { + cancelled = true; + }; + }, [open, refreshKey]); + + const selectedModel = useMemo( + () => models.find((model) => model.id === value) ?? null, + [models, value], + ); + + const filteredModels = useMemo(() => { + const needle = normalizeForSearch(query.trim()); + if (!needle) { + return models; + } + return models.filter((model) => { + const haystack = normalizeForSearch( + `${model.id} ${model.display_name} ${model.model_id ?? ""} ${model.path}`, + ); + return haystack.includes(needle); + }); + }, [models, query]); + + const selectModel = useCallback( + async (model: LocalModelInfo) => { + if (isExpandableGguf(model)) { + setExpandedModelId((current) => + current === model.id ? null : model.id, + ); + return; + } + if (!isDirectGguf(model)) { + setProbingVariantModelId(model.id); + try { + const response = await listGgufVariants(model.id); + if (response.variants.length > 0) { + setExpandedModelId(model.id); + return; + } + } catch { + // Non-GGUF local models commonly have no variant endpoint. Fall + // through to regular selection so users can still choose them. + } finally { + setProbingVariantModelId(null); + } + } + onChange(model.id, null); + setOpen(false); + }, + [onChange], + ); + + const selectVariant = useCallback( + (modelId: string, variant: string) => { + onChange(modelId, variant); + setOpen(false); + }, + [onChange], + ); + + return ( + + + + + +
+
+
+ setQuery(event.target.value)} + placeholder="Filter local models" + className="h-8 flex-1" + autoFocus={true} + /> + +
+
+ +
event.stopPropagation()} + > + +
+
+ + + ); +} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx index 368ae08acb..68f912bc57 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx @@ -1,12 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import { Checkbox } from "@/components/ui/checkbox"; import { Collapsible, CollapsibleContent, CollapsibleTrigger, } from "@/components/ui/collapsible"; -import { Checkbox } from "@/components/ui/checkbox"; import { Combobox, ComboboxContent, @@ -22,6 +22,7 @@ import type { ModelConfig } from "../../types"; import { CollapsibleSectionTriggerButton } from "../shared/collapsible-section-trigger"; import { FieldLabel } from "../shared/field-label"; import { NameField } from "../shared/name-field"; +import { LocalRecipeModelSelector } from "./local-recipe-model-selector"; type ModelConfigDialogProps = { config: ModelConfig; @@ -45,6 +46,7 @@ export function ModelConfigDialog({ const maxTokensId = `${config.id}-max-tokens`; const timeoutId = `${config.id}-timeout`; const extraBodyId = `${config.id}-inference-extra-body`; + const skipHealthCheckId = `${config.id}-skip-health-check`; const providerAnchorRef = useRef(null); const providerInputRef = useRef(config.provider); // Sync providerInputRef with the current provider value. Updating a ref in @@ -61,16 +63,25 @@ export function ModelConfigDialog({ onUpdate({ [key]: value } as Partial); }; - // Apply provider selection while keeping the local-provider model autofill - // consistent across both dropdown selection and free-typed + blur input. + // Apply provider selection while clearing model identifiers that only make + // sense for the previous provider locality. const applyProviderChange = (selectedProvider: string) => { - const isLocal = localProviderNames.has(selectedProvider); - if (isLocal && !config.model.trim()) { - onUpdate({ provider: selectedProvider, model: "local" }); + const nextIsLocal = localProviderNames.has(selectedProvider); + if (isLinkedToLocal !== nextIsLocal) { + onUpdate({ + provider: selectedProvider, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }); return; } - if (!isLocal && config.model === "local") { - onUpdate({ provider: selectedProvider, model: "" }); + if (!nextIsLocal) { + onUpdate({ + provider: selectedProvider, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }); return; } updateField("provider", selectedProvider); @@ -88,8 +99,8 @@ export function ModelConfigDialog({ Set up one reusable model choice for your AI steps

- Choose the provider connection, enter the exact model ID, then save any - generation defaults you want to reuse. + Choose the provider connection, enter the exact model ID, then save + any generation defaults you want to reuse.

@@ -144,15 +155,48 @@ export function ModelConfigDialog({ - updateField("model", event.target.value)} + hint={ + isLinkedToLocal + ? "Choose the local model Recipes should load before Run or Validate." + : "The exact model name sent to the connection." + } /> + {isLinkedToLocal ? ( + + onUpdate({ + model, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: variant ?? undefined, + }) + } + /> + ) : ( + + onUpdate({ + model: event.target.value, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }) + } + /> + )} + {isLinkedToLocal ? ( +

+ Recipes will load this model automatically. GGUF quantization is + saved with the preset. +

+ ) : null}
@@ -250,8 +294,12 @@ export function ModelConfigDialog({ } />
-
@@ -149,15 +187,28 @@ export function GithubCrawlerEasyView({
- handleModelChange(event.target.value)} - placeholder="unsloth/gemma-4-E2B-it-GGUF" - disabled={!modelConfig} + hint={ + isModelLinkedToLocal + ? "Choose the local model this recipe should load." + : "OpenAI-compatible model id." + } /> + {isModelLinkedToLocal ? ( + + ) : ( + handleModelChange(event.target.value)} + placeholder="unsloth/gemma-4-E2B-it-GGUF" + disabled={!modelConfig} + /> + )}
diff --git a/studio/frontend/src/features/recipe-studio/executions/tracker.ts b/studio/frontend/src/features/recipe-studio/executions/tracker.ts index d83f662fcf..97e70c66c4 100644 --- a/studio/frontend/src/features/recipe-studio/executions/tracker.ts +++ b/studio/frontend/src/features/recipe-studio/executions/tracker.ts @@ -40,6 +40,11 @@ type TrackRecipeExecutionParams = { onPreviewSuccess?: () => void; }; +export type TrackRecipeExecutionResult = { + success: boolean; + terminal: boolean; +}; + function isTerminalStatus(status: RecipeExecutionStatus): boolean { return status === "completed" || status === "error" || status === "cancelled"; } @@ -53,7 +58,8 @@ function normalizeCompletedProgress(input: { } { const { latestExecution, rows } = input; const progressTotal = - typeof latestExecution.progress?.total === "number" && latestExecution.progress.total > 0 + typeof latestExecution.progress?.total === "number" && + latestExecution.progress.total > 0 ? latestExecution.progress.total : latestExecution.rows > 0 ? latestExecution.rows @@ -92,7 +98,7 @@ export async function trackRecipeExecution({ onUpsert, onSetPreviewErrors, onPreviewSuccess, -}: TrackRecipeExecutionParams): Promise { +}: TrackRecipeExecutionParams): Promise { let done = false; let lastStatus: RecipeExecutionStatus = initialExecution.status; let completedEventPayload: Record | null = null; @@ -124,7 +130,9 @@ export async function trackRecipeExecution({ } const eventType = - typeof event.payload.type === "string" ? event.payload.type : event.event; + typeof event.payload.type === "string" + ? event.payload.type + : event.event; if (eventType === "job.started") { latestExecution = { @@ -163,7 +171,7 @@ export async function trackRecipeExecution({ error: typeof event.payload.error === "string" ? event.payload.error - : latestExecution.error ?? `${label} failed.`, + : (latestExecution.error ?? `${label} failed.`), }; onUpsert(latestExecution); return; @@ -178,6 +186,19 @@ export async function trackRecipeExecution({ return; } + if (eventType === "job.cancelled") { + lastStatus = "cancelled"; + done = true; + latestExecution = { + ...latestExecution, + status: "cancelled", + finishedAt: Date.now(), + error: latestExecution.error ?? "Run cancelled.", + }; + onUpsert(latestExecution); + return; + } + if (changed) { onUpsert(latestExecution); } @@ -189,6 +210,9 @@ export async function trackRecipeExecution({ try { while (!done) { const status = await getRecipeJobStatus(jobId); + if (done && isTerminalStatus(lastStatus)) { + break; + } const mappedStatus = mapJobStatus(status.status); lastStatus = mappedStatus; latestExecution = applyExecutionStatusSnapshot(latestExecution, status); @@ -200,18 +224,19 @@ export async function trackRecipeExecution({ } } } catch (error) { - const message = toErrorMessage(error, `${label} failed.`); - latestExecution = { - ...latestExecution, - status: "error", - error: message, - finishedAt: Date.now(), - }; - onUpsert(latestExecution); - if (notify) { - toastError(`${label} failed`, message); + const terminal = isTerminalStatus(lastStatus); + if (!terminal) { + const message = toErrorMessage(error, `${label} failed.`); + latestExecution = { + ...latestExecution, + error: message, + }; + onUpsert(latestExecution); + if (notify) { + toastError(`${label} failed`, message); + } + return { success: false, terminal: false }; } - return false; } finally { eventsAbortController.abort(); } @@ -220,7 +245,10 @@ export async function trackRecipeExecution({ for (let attempt = 0; attempt < 3; attempt += 1) { try { const finalStatus = await getRecipeJobStatus(jobId); - latestExecution = applyExecutionStatusSnapshot(latestExecution, finalStatus); + latestExecution = applyExecutionStatusSnapshot( + latestExecution, + finalStatus, + ); } catch { break; } @@ -229,19 +257,20 @@ export async function trackRecipeExecution({ } } - const eventAnalysis = completedEventPayload - ? completedEventPayload["analysis"] - : null; - const eventDataset = completedEventPayload - ? completedEventPayload["dataset"] - : null; + const completedPayload = completedEventPayload as Record< + string, + unknown + > | null; + const eventAnalysis = completedPayload ? completedPayload.analysis : null; + const eventDataset = completedPayload ? completedPayload.dataset : null; const eventProcessorArtifacts = - completedEventPayload && - typeof completedEventPayload["processor_artifacts"] === "object" && - completedEventPayload["processor_artifacts"] !== null - ? (completedEventPayload["processor_artifacts"] as Record) + completedPayload && + typeof completedPayload.processor_artifacts === "object" && + completedPayload.processor_artifacts !== null + ? (completedPayload.processor_artifacts as Record) : null; - const shouldFetchPreviewDataset = kind === "preview" && !Array.isArray(eventDataset); + const shouldFetchPreviewDataset = + kind === "preview" && !Array.isArray(eventDataset); const shouldFetchAnalysis = !completedEventPayload || typeof eventAnalysis !== "object" || @@ -262,9 +291,7 @@ export async function trackRecipeExecution({ ? normalizeAnalysis(analysisResult.value) : latestExecution.analysis; const datasetResponse = - datasetResult.status === "fulfilled" - ? datasetResult.value - : null; + datasetResult.status === "fulfilled" ? datasetResult.value : null; const dataset = datasetResponse ? normalizeDatasetRows(datasetResponse.dataset) : latestExecution.dataset; @@ -272,7 +299,10 @@ export async function trackRecipeExecution({ datasetResponse && typeof datasetResponse.total === "number" ? datasetResponse.total : latestExecution.datasetTotal; - const completedProgress = normalizeCompletedProgress({ latestExecution, rows }); + const completedProgress = normalizeCompletedProgress({ + latestExecution, + rows, + }); latestExecution = { ...latestExecution, @@ -285,7 +315,8 @@ export async function trackRecipeExecution({ datasetPage: 1, datasetPageSize: DATASET_PAGE_SIZE, error: null, - processor_artifacts: eventProcessorArtifacts ?? latestExecution.processor_artifacts, + processor_artifacts: + eventProcessorArtifacts ?? latestExecution.processor_artifacts, finishedAt: latestExecution.finishedAt ?? Date.now(), }; onUpsert(latestExecution); @@ -299,7 +330,7 @@ export async function trackRecipeExecution({ toastSuccess("Full run completed."); } } - return true; + return { success: true, terminal: true }; } if (lastStatus === "cancelled") { @@ -313,7 +344,7 @@ export async function trackRecipeExecution({ if (notify) { toastError(`${label} cancelled`, "The execution was cancelled."); } - return false; + return { success: false, terminal: true }; } latestExecution = { @@ -326,5 +357,5 @@ export async function trackRecipeExecution({ if (notify) { toastError(`${label} failed`, latestExecution.error ?? "Execution failed."); } - return false; + return { success: false, terminal: true }; } diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts index c3da5b1999..19a6a3004d 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts @@ -1,14 +1,11 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { useCallback, useEffect, useState } from "react"; -import { useShallow } from "zustand/react/shallow"; +import { getInferenceStatus, loadModel } from "@/features/chat"; import { toast } from "@/lib/toast"; import { toastError } from "@/shared/toast"; -import { - getInferenceStatus, - loadModel, -} from "@/features/chat/api/chat-api"; +import { useCallback, useEffect, useState } from "react"; +import { useShallow } from "zustand/react/shallow"; import { cancelRecipeJob, createRecipeJob, @@ -23,8 +20,8 @@ import type { import { DATASET_PAGE_SIZE, executionLabel, - normalizeRunName, normalizeDatasetRows, + normalizeRunName, toErrorMessage, withExecutionDefaults, } from "../executions/execution-helpers"; @@ -32,84 +29,243 @@ import { findResumableExecution, loadSortedRecipeExecutions, } from "../executions/hydration"; -import { createBaseExecutionRecord } from "../executions/runtime"; import { buildExecutionPayload, sanitizeExecutionRows, } from "../executions/run-settings"; +import { createBaseExecutionRecord } from "../executions/runtime"; import { trackRecipeExecution } from "../executions/tracker"; import { type RecipeRunSettings, useRecipeExecutionsStore, } from "../stores/recipe-executions"; -import type { RecipePayload, RecipePayloadResult } from "../utils/payload/types"; +import type { + RecipePayload, + RecipePayloadResult, +} from "../utils/payload/types"; -/** - * Auto-load the local model before running a recipe that uses it. - * - * Looks at payload.recipe.model_providers for any provider with is_local=true, - * finds the bound model_configs and asks the backend to load whichever model - * the first local-bound model_config points at. Skips when the inference - * server already has that exact model active. This removes the "open /chat - * first" prerequisite that users kept tripping on. - */ -async function ensureLocalModelLoaded( - payload: RecipePayload, -): Promise { +const GGUF_MODEL_PATTERN = /gguf/i; + +function collectUsedLlmModelAliases(payload: RecipePayload): Set { + const columns = Array.isArray(payload.recipe.columns) + ? payload.recipe.columns + : []; + const aliases = new Set(); + for (const column of columns) { + const columnType = column.column_type; + if (typeof columnType !== "string" || !columnType.startsWith("llm-")) { + continue; + } + const alias = column.model_alias; + if (typeof alias === "string" && alias.trim()) { + aliases.add(alias.trim()); + } + } + return aliases; +} + +type LocalModelSelection = { + target: string; + ggufVariant: string; + aliases: string[]; +}; + +type LocalModelLoadPlan = + | { selection: LocalModelSelection; error: null; legacyAliases?: never } + | { selection: null; error: string; legacyAliases?: never } + | { selection: null; error: null; legacyAliases: string[] }; + +type RestorableLocalModelSnapshot = { + selection: LocalModelSelection | null; + unrestorableLabel: string | null; +}; + +function getLocalProviderNames(payload: RecipePayload): Set { const providers = Array.isArray(payload.recipe.model_providers) - ? (payload.recipe.model_providers as Array>) + ? (payload.recipe.model_providers as Record[]) : []; const localProviderNames = new Set(); - for (const p of providers) { - if (p.is_local === true && typeof p.name === "string") { - localProviderNames.add(p.name); + for (const provider of providers) { + if (provider.is_local === true && typeof provider.name === "string") { + localProviderNames.add(provider.name); } } - if (localProviderNames.size === 0) { - return null; + return localProviderNames; +} + +function findUsedLocalModelConfigs( + payload: RecipePayload, + localProviderNames: Set, +): Record[] { + const usedAliases = collectUsedLlmModelAliases(payload); + if (usedAliases.size === 0) { + return []; } const modelConfigs = Array.isArray(payload.recipe.model_configs) - ? (payload.recipe.model_configs as Array>) + ? payload.recipe.model_configs : []; - const boundConfig = modelConfigs.find( - (c) => typeof c.provider === "string" && localProviderNames.has(c.provider), - ); + return modelConfigs.filter((config) => { + const provider = config.provider; + const alias = config.alias; + return ( + typeof provider === "string" && + localProviderNames.has(provider) && + typeof alias === "string" && + usedAliases.has(alias) + ); + }); +} + +function readLocalModelSelection( + boundConfig: Record, +): LocalModelLoadPlan { + const alias = + typeof boundConfig.alias === "string" ? boundConfig.alias : "local model"; const target = - typeof boundConfig?.model === "string" ? boundConfig.model.trim() : ""; + typeof boundConfig.model === "string" ? boundConfig.model.trim() : ""; + const ggufVariant = + typeof boundConfig.gguf_variant === "string" + ? boundConfig.gguf_variant.trim() + : ""; if (!target) { - return null; + return { + selection: null, + error: `Model config ${alias}: choose a local model before validating or running this recipe.`, + }; + } + if (target.toLowerCase() === "local") { + return { selection: null, error: null, legacyAliases: [alias] }; + } + return { selection: { target, ggufVariant, aliases: [alias] }, error: null }; +} + +function getLocalModelLoadPlan( + boundConfigs: Record[], +): LocalModelLoadPlan | null { + const selections = new Map(); + const legacyAliases: string[] = []; + for (const boundConfig of boundConfigs) { + const next = readLocalModelSelection(boundConfig); + if (next.error) { + return next; + } + if (next.legacyAliases) { + legacyAliases.push(...next.legacyAliases); + continue; + } + const selection = next.selection; + if (!selection) { + continue; + } + const key = `${selection.target.toLowerCase()}\u0000${selection.ggufVariant}`; + const existing = selections.get(key); + if (existing) { + existing.aliases.push(...selection.aliases); + continue; + } + selections.set(key, selection); } + if (legacyAliases.length > 0 && selections.size > 0) { + const aliases = [ + ...legacyAliases, + ...[...selections.values()].flatMap((selection) => selection.aliases), + ].join(", "); + return { + selection: null, + error: `Recipes found mixed legacy and selected local models. Reselect the same concrete local model for: ${aliases}.`, + }; + } + + if (legacyAliases.length > 0) { + return { selection: null, error: null, legacyAliases }; + } + + if (selections.size > 1) { + const aliases = [...selections.values()] + .flatMap((selection) => selection.aliases) + .join(", "); + return { + selection: null, + error: `Recipes supports one active local model per run. Select the same local model and GGUF variant for: ${aliases}.`, + }; + } + + const selection = [...selections.values()][0]; + return selection ? { selection, error: null } : null; +} + +function isDirectGgufTarget(target: string): boolean { + return target.toLowerCase().endsWith(".gguf"); +} + +function localSelectionMatchesActive(input: { + target: string; + ggufVariant: string; + activeModel: string | null | undefined; + activeVariant: string; +}): boolean { + const { target, ggufVariant, activeModel, activeVariant } = input; + if (!activeModel || activeModel.toLowerCase() !== target.toLowerCase()) { + return false; + } + return ( + activeVariant === ggufVariant || + (isDirectGgufTarget(target) && !ggufVariant) + ); +} + +async function isLocalModelAlreadyLoaded( + selection: LocalModelSelection, +): Promise { + const { target, ggufVariant } = selection; try { const status = await getInferenceStatus(); - if ( - status.active_model && - status.active_model.toLowerCase() === target.toLowerCase() - ) { - return null; - } + return localSelectionMatchesActive({ + target, + ggufVariant, + activeModel: status.model_identifier ?? status.active_model, + activeVariant: status.gguf_variant?.trim() ?? "", + }); } catch { // Fall through to load attempt; the backend will re-error if needed. + return false; } +} - const toastId = toast.loading(`Loading ${target}…`, { +async function loadLocalModelSelection( + selection: LocalModelSelection, +): Promise { + const { target, ggufVariant } = selection; + const modelLabel = ggufVariant ? `${target} (${ggufVariant})` : target; + const toastId = toast.loading(`Loading ${modelLabel}...`, { description: "Starting the local inference server for this recipe.", }); try { - const isGguf = /gguf/i.test(target); + const isGguf = GGUF_MODEL_PATTERN.test(target) || Boolean(ggufVariant); await loadModel({ + // biome-ignore lint/style/useNamingConvention: api schema model_path: target, + // biome-ignore lint/style/useNamingConvention: api schema hf_token: null, + // biome-ignore lint/style/useNamingConvention: api schema max_seq_length: isGguf ? 0 : 4096, + // biome-ignore lint/style/useNamingConvention: api schema load_in_4bit: true, + // biome-ignore lint/style/useNamingConvention: api schema is_lora: false, - gguf_variant: null, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: ggufVariant || null, + // biome-ignore lint/style/useNamingConvention: api schema trust_remote_code: false, + // biome-ignore lint/style/useNamingConvention: api schema chat_template_override: null, + // biome-ignore lint/style/useNamingConvention: api schema cache_type_kv: null, + // biome-ignore lint/style/useNamingConvention: api schema speculative_type: null, }); - toast.success(`Loaded ${target}`, { id: toastId, duration: 2000 }); + toast.success(`Loaded ${modelLabel}`, { id: toastId, duration: 2000 }); return null; } catch (error) { toast.dismiss(toastId); @@ -117,6 +273,147 @@ async function ensureLocalModelLoaded( } } +function getLocalModelLoadPlanForPayload( + payload: RecipePayload, +): LocalModelLoadPlan | null { + const localProviderNames = getLocalProviderNames(payload); + if (localProviderNames.size === 0) { + return null; + } + + const boundConfigs = findUsedLocalModelConfigs(payload, localProviderNames); + return getLocalModelLoadPlan(boundConfigs); +} + +async function getActiveLocalModelSelection(): Promise { + try { + const status = await getInferenceStatus(); + const target = status.active_model?.trim(); + if (!target) { + return null; + } + return { + target, + ggufVariant: status.gguf_variant?.trim() ?? "", + aliases: ["previous Chat model"], + }; + } catch { + return null; + } +} + +async function getRestorableActiveLocalModelSelection(): Promise { + try { + const status = await getInferenceStatus(); + const activeLabel = status.active_model?.trim() ?? null; + const target = ( + status.model_identifier ?? (status.is_gguf ? null : status.active_model) + )?.trim(); + if (!target) { + return { + selection: null, + unrestorableLabel: activeLabel, + }; + } + return { + selection: { + target, + ggufVariant: status.gguf_variant?.trim() ?? "", + aliases: ["previous Chat model"], + }, + unrestorableLabel: null, + }; + } catch { + return { selection: null, unrestorableLabel: null }; + } +} + +function isSameLocalModelSelection( + left: LocalModelSelection | null, + right: LocalModelSelection, +): boolean { + return Boolean( + left && + left.target.toLowerCase() === right.target.toLowerCase() && + left.ggufVariant === right.ggufVariant, + ); +} + +async function ensureLocalModelLoaded( + payload: RecipePayload, +): Promise { + const loadPlan = getLocalModelLoadPlanForPayload(payload); + if (!loadPlan) { + return null; + } + if (loadPlan.legacyAliases) { + const activeSelection = await getActiveLocalModelSelection(); + return activeSelection + ? null + : `Existing recipe uses legacy local model for ${loadPlan.legacyAliases.join(", ")}. Select a concrete local model or load one in Chat.`; + } + if (!loadPlan.selection) { + return loadPlan.error; + } + if (await isLocalModelAlreadyLoaded(loadPlan.selection)) { + return null; + } + return loadLocalModelSelection(loadPlan.selection); +} + +async function prepareLocalModelForRun(payload: RecipePayload): Promise<{ + error: string | null; + restorePrevious: (() => Promise) | null; +}> { + const loadPlan = getLocalModelLoadPlanForPayload(payload); + if (!loadPlan) { + return { error: null, restorePrevious: null }; + } + if (loadPlan.legacyAliases) { + const activeSelection = await getActiveLocalModelSelection(); + return activeSelection + ? { error: null, restorePrevious: null } + : { + error: `Existing recipe uses legacy local model for ${loadPlan.legacyAliases.join(", ")}. Select a concrete local model or load one in Chat.`, + restorePrevious: null, + }; + } + if (!loadPlan.selection) { + return { error: loadPlan.error, restorePrevious: null }; + } + if (await isLocalModelAlreadyLoaded(loadPlan.selection)) { + return { error: null, restorePrevious: null }; + } + + const previousSnapshot = await getRestorableActiveLocalModelSelection(); + const previousSelection = previousSnapshot.selection; + const error = await loadLocalModelSelection(loadPlan.selection); + if (error) { + return { error, restorePrevious: null }; + } + if (isSameLocalModelSelection(previousSelection, loadPlan.selection)) { + return { error: null, restorePrevious: null }; + } + return { + error: null, + restorePrevious: previousSelection + ? async () => { + const restoreError = await loadLocalModelSelection(previousSelection); + if (restoreError) { + toastError("Could not restore previous local model", restoreError); + } + } + : previousSnapshot.unrestorableLabel + ? () => { + toast.warning("Previous local model was not restored", { + description: `${previousSnapshot.unrestorableLabel} was selected from a native file path. Reopen it in Chat to continue with that model.`, + }); + return Promise.resolve(); + } + : null, + }; +} + type UseRecipeExecutionsParams = { recipeId: string; currentSignature: string; @@ -161,7 +458,11 @@ type UseRecipeExecutionsResult = { }; function formatValidationMessages(input: { - errors: Array<{ message: string; path?: string | null; code?: string | null }>; + errors: Array<{ + message: string; + path?: string | null; + code?: string | null; + }>; }): string[] { return input.errors.map((item) => { const path = item.path?.trim(); @@ -249,7 +550,8 @@ export function useRecipeExecutions({ (record: RecipeExecutionRecord): void => { const normalizedRecord = withExecutionDefaults(record); upsertExecution(normalizedRecord); - void saveRecipeExecution(normalizedRecord).catch((error) => { + saveRecipeExecution(normalizedRecord).catch((error) => { + // biome-ignore lint/suspicious/noConsole: background persistence failures should not interrupt the UI console.error("Save recipe execution failed:", error); }); }, @@ -287,7 +589,7 @@ export function useRecipeExecutions({ return; } - void trackRecipeExecution({ + trackRecipeExecution({ label: executionLabel(resumable.kind), kind: resumable.kind, rows: resumable.rows, @@ -299,11 +601,12 @@ export function useRecipeExecutions({ onPreviewSuccess, }); } catch (error) { + // biome-ignore lint/suspicious/noConsole: hydration failures are non-blocking diagnostics console.error("Load recipe executions failed:", error); } } - void hydrate(); + hydrate(); return () => { cancelled = true; @@ -344,9 +647,11 @@ export function useRecipeExecutions({ rows: number; settings: RecipeRunSettings; runName: string | null; + restorePrevious?: (() => Promise) | null; }): Promise => { - const { kind, payload, rows, settings, runName } = input; - const setLoading = kind === "preview" ? setPreviewLoading : setFullLoading; + const { kind, payload, rows, settings, runName, restorePrevious } = input; + const setLoading = + kind === "preview" ? setPreviewLoading : setFullLoading; const label = executionLabel(kind); setLoading(true); @@ -362,6 +667,8 @@ export function useRecipeExecutions({ onExecutionStart?.(); setRunDialogOpen(false); + let jobCreated = false; + let shouldRestorePrevious = false; try { const jobPayload = buildExecutionPayload({ payload, @@ -371,13 +678,14 @@ export function useRecipeExecutions({ runName, }); const createdJob = await createRecipeJob(jobPayload); + jobCreated = true; const executionWithJob = { ...baseExecution, jobId: createdJob.job_id, }; upsertAndPersist(executionWithJob); - return await trackRecipeExecution({ + const tracked = await trackRecipeExecution({ label, kind, rows, @@ -388,6 +696,8 @@ export function useRecipeExecutions({ onSetPreviewErrors: setRunErrors, onPreviewSuccess, }); + shouldRestorePrevious = tracked.terminal; + return tracked.success; } catch (error) { const message = toErrorMessage(error, `${label} request failed.`); upsertAndPersist({ @@ -398,8 +708,14 @@ export function useRecipeExecutions({ }); setRunErrors([message]); toastError(`${label} failed`, message); + if (!jobCreated) { + shouldRestorePrevious = true; + } return false; } finally { + if (shouldRestorePrevious && restorePrevious) { + await restorePrevious(); + } setLoading(false); } }, @@ -416,6 +732,48 @@ export function useRecipeExecutions({ ], ); + const prepareLocalModelForExecution = useCallback( + async ( + payload: RecipePayload, + ): Promise<(() => Promise) | null | false> => { + const { error, restorePrevious } = await prepareLocalModelForRun(payload); + if (!error) { + return restorePrevious; + } + setRunErrors([error]); + toastError("Local model failed to load", error); + return false; + }, + [setRunErrors], + ); + + const validateExecutionPayload = useCallback( + async ( + executionPayload: Parameters[0], + ): Promise => { + try { + const validation = await validateRecipe(executionPayload); + if (validation.valid) { + return true; + } + const errors = formatValidationMessages({ + errors: validation.errors, + }); + const fallback = validation.raw_detail ?? "Validation failed."; + const nextErrors = errors.length > 0 ? errors : [fallback]; + setRunErrors(nextErrors); + toastError("Validation failed", nextErrors[0]); + return false; + } catch (error) { + const message = toErrorMessage(error, "Validation failed."); + setRunErrors([message]); + toastError("Validation failed", message); + return false; + } + }, + [setRunErrors], + ); + const runWithValidation = useCallback( async ( kind: RecipeExecutionKind, @@ -435,20 +793,11 @@ export function useRecipeExecutions({ return false; } - // Flip to the Runs pane BEFORE we run ensureLocalModelLoaded + validate. - // Validation re-crawls the seed (multiple seconds for the github_repo - // reader) and the user otherwise stares at a "Running..." button with - // nothing else changing. runExecution() later no-ops this callback if - // the view has already been flipped, so we fire it once here. + // Flip to the Runs pane before validation starts. Validation can re-crawl + // the seed (multiple seconds for the github_repo reader), and runExecution() + // later no-ops this callback if the view has already been flipped. onExecutionStart?.(); - const localLoadError = await ensureLocalModelLoaded(payload); - if (localLoadError) { - setRunErrors([localLoadError]); - toastError("Local model failed to load", localLoadError); - return false; - } - const normalizedRows = sanitizeExecutionRows(rows, kind); const executionPayload = buildExecutionPayload({ payload, @@ -458,20 +807,17 @@ export function useRecipeExecutions({ runName, }); - try { - const validation = await validateRecipe(executionPayload); - if (!validation.valid) { - const errors = formatValidationMessages({ errors: validation.errors }); - const fallback = validation.raw_detail ?? "Validation failed."; - const nextErrors = errors.length > 0 ? errors : [fallback]; - setRunErrors(nextErrors); - toastError("Validation failed", nextErrors[0]); - return false; - } - } catch (error) { - const message = toErrorMessage(error, "Validation failed."); - setRunErrors([message]); - toastError("Validation failed", message); + if (!(await validateExecutionPayload(executionPayload))) { + return false; + } + + // Recipe and Chat share one singleton local inference backend. This + // direct load is a point-in-time handoff to job creation, not a lease: + // if Chat swaps models after this succeeds, the backend will reject or + // run against the active backend state. A future generation token should + // be validated across this load and the `/jobs` loaded-model gate. + const restorePrevious = await prepareLocalModelForExecution(payload); + if (restorePrevious === false) { return false; } @@ -481,26 +827,29 @@ export function useRecipeExecutions({ rows: normalizedRows, settings: runSettings, runName, + restorePrevious, }); }, [ onExecutionStart, + prepareLocalModelForExecution, readExecutablePayload, runExecution, runSettings, setRunErrors, + validateExecutionPayload, ], ); - const runPreview = useCallback(async (): Promise => { + const runPreview = useCallback((): Promise => { return runWithValidation("preview", previewRows, null); }, [previewRows, runWithValidation]); - const runFull = useCallback(async (): Promise => { + const runFull = useCallback((): Promise => { return runWithValidation("full", fullRows, fullRunName); }, [fullRows, fullRunName, runWithValidation]); - const runFromDialog = useCallback(async (): Promise => { + const runFromDialog = useCallback((): Promise => { setValidateResult(null); if (runDialogKind === "preview") { return runPreview(); @@ -512,9 +861,10 @@ export function useRecipeExecutions({ setRunErrors([]); const payload = readPayload(); if (!payload) { - const nextErrors = payloadResult.errors.length > 0 - ? payloadResult.errors - : [payloadErrorMessage]; + const nextErrors = + payloadResult.errors.length > 0 + ? payloadResult.errors + : [payloadErrorMessage]; setValidateResult({ valid: false, errors: nextErrors, @@ -525,24 +875,46 @@ export function useRecipeExecutions({ const rows = runDialogKind === "preview" ? previewRows : fullRows; const normalizedRows = sanitizeExecutionRows(rows, runDialogKind); - const executionPayload = buildExecutionPayload({ - payload, - kind: runDialogKind, - rows: normalizedRows, - settings: runSettings, - runName: runDialogKind === "full" ? normalizeRunName(fullRunName) : null, - }); setValidateLoading(true); try { + const executionPayload = buildExecutionPayload({ + payload, + kind: runDialogKind, + rows: normalizedRows, + settings: runSettings, + runName: + runDialogKind === "full" ? normalizeRunName(fullRunName) : null, + }); const validation = await validateRecipe(executionPayload); const errors = formatValidationMessages({ errors: validation.errors }); + if (!validation.valid) { + setValidateResult({ + valid: false, + errors, + rawDetail: validation.raw_detail ?? null, + }); + return false; + } + + const localLoadError = await ensureLocalModelLoaded(payload); + if (localLoadError) { + setRunErrors([localLoadError]); + setValidateResult({ + valid: false, + errors: [localLoadError], + rawDetail: null, + }); + toastError("Local model failed to load", localLoadError); + return false; + } + setValidateResult({ - valid: validation.valid, + valid: true, errors, rawDetail: validation.raw_detail ?? null, }); - return validation.valid; + return true; } catch (error) { const message = toErrorMessage(error, "Validation failed."); setValidateResult({ @@ -612,7 +984,12 @@ export function useRecipeExecutions({ const loadExecutionDatasetPage = useCallback( async (id: string, page: number): Promise => { const execution = executions.find((entry) => entry.id === id); - if (!execution || execution.kind !== "full" || !execution.jobId || page < 1) { + if ( + !execution || + execution.kind !== "full" || + !execution.jobId || + page < 1 + ) { return; } @@ -625,7 +1002,9 @@ export function useRecipeExecutions({ }); const dataset = normalizeDatasetRows(response.dataset); const total = - typeof response.total === "number" ? response.total : execution.datasetTotal; + typeof response.total === "number" + ? response.total + : execution.datasetTotal; upsertAndPersist({ ...execution, dataset, diff --git a/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts b/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts index 56634b84c8..dd5e12233c 100644 --- a/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts +++ b/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts @@ -97,7 +97,9 @@ export function applyRenameToConfig( next = { ...base, // biome-ignore lint/style/useNamingConvention: api schema - target_columns: targets.map((target) => (target === from ? to : target)), + target_columns: targets.map((target) => + target === from ? to : target, + ), }; } } @@ -137,14 +139,12 @@ export function applyRemovalToConfig( } if (config.kind === "model_config" && config.provider === ref) { const base = next as ModelConfig; - // Clear the synthetic "local" placeholder when the provider that was - // a local provider is removed; otherwise the stale placeholder would - // pass validation against a future external provider and then fail - // at runtime against a real API ("model not found"). next = { ...base, provider: "", - model: base.model === "local" ? "" : base.model, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, }; } if (config.kind === "llm" && config.model_alias === ref) { @@ -156,7 +156,9 @@ export function applyRemovalToConfig( next = { ...base, tool_alias: "" }; } if (config.kind === "validator") { - const targets = (config.target_columns ?? []).filter((target) => target !== ref); + const targets = (config.target_columns ?? []).filter( + (target) => target !== ref, + ); if (targets.length !== (config.target_columns ?? []).length) { const base = next as typeof config; next = { @@ -206,5 +208,7 @@ export function applyRemovalToConfigs( if (!ref) { return configs; } - return applyConfigTransform(configs, (config) => applyRemovalToConfig(config, ref)); + return applyConfigTransform(configs, (config) => + applyRemovalToConfig(config, ref), + ); } diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 036cc94082..52ce1329bf 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -12,23 +12,23 @@ import { applyNodeChanges, } from "@xyflow/react"; import { create } from "zustand"; -import type { - RecipeNode, - RecipeProcessorConfig, - LayoutDirection, - LlmType, - NodeConfig, - SeedSourceType, - SamplerType, -} from "../types"; import { - getBlockDefinition, type BlockKind, type BlockType, type SeedBlockType, + getBlockDefinition, } from "../blocks/registry"; -import { deriveDisplayGraph } from "../utils/graph/derive-display-graph"; +import type { + LayoutDirection, + LlmType, + NodeConfig, + RecipeNode, + RecipeProcessorConfig, + SamplerType, + SeedSourceType, +} from "../types"; import { applyRecipeConnection, isValidRecipeConnection } from "../utils/graph"; +import { deriveDisplayGraph } from "../utils/graph/derive-display-graph"; import { HANDLE_IDS, normalizeRecipeHandleId, @@ -42,8 +42,8 @@ import { } from "./helpers/model-infra-layout"; import { applyEdgeRemovals, applyNodeRemovals } from "./helpers/removals"; import { - applyRenameToConfigs, applyLayoutDirectionToNodes, + applyRenameToConfigs, buildNodeUpdate, syncEdgesForConfigPatch, syncSubcategoryConfigsForCategoryUpdate, @@ -97,7 +97,11 @@ type RecipeStudioState = { position?: XYPosition, openDialog?: boolean, ) => void; - addLlmNode: (type: LlmType, position?: XYPosition, openDialog?: boolean) => void; + addLlmNode: ( + type: LlmType, + position?: XYPosition, + openDialog?: boolean, + ) => void; addModelProviderNode: (position?: XYPosition, openDialog?: boolean) => void; addModelConfigNode: (position?: XYPosition, openDialog?: boolean) => void; addToolProfileNode: (position?: XYPosition, openDialog?: boolean) => void; @@ -250,7 +254,10 @@ function connectSemantic( }; } -function isModelSemanticEdge(edge: Edge, configs: Record): boolean { +function isModelSemanticEdge( + edge: Edge, + configs: Record, +): boolean { const source = configs[edge.source]; const target = configs[edge.target]; return Boolean( @@ -315,12 +322,16 @@ export const useRecipeStudioStore = create((set, get) => ({ auxNodePositions: {}, llmAuxVisibility: state.llmAuxVisibility, }); - const { nodes } = getLayoutedElements(displayGraph.nodes, displayGraph.edges, { - direction: state.layoutDirection, - nodesep: isTopBottom ? 120 : 80, - ranksep: isTopBottom ? 140 : 80, - configs: state.configs, - }); + const { nodes } = getLayoutedElements( + displayGraph.nodes, + displayGraph.edges, + { + direction: state.layoutDirection, + nodesep: isTopBottom ? 120 : 80, + ranksep: isTopBottom ? 140 : 80, + configs: state.configs, + }, + ); const layoutedPositions = new Map( nodes.map((node) => [node.id, node.position] as const), ); @@ -381,13 +392,7 @@ export const useRecipeStudioStore = create((set, get) => ({ (config) => config.kind === "seed", ); if (!existing) { - return buildAddedNodeState( - state, - "seed", - type, - position, - openDialog, - ); + return buildAddedNodeState(state, "seed", type, position, openDialog); } let nextSourceType: SeedSourceType = "hf"; if (type === "seed_local") { @@ -430,7 +435,10 @@ export const useRecipeStudioStore = create((set, get) => ({ [existing.id]: nextConfig, }, nodes: updateNodeData( - state.nodes.map((node) => ({ ...node, selected: node.id === existing.id })), + state.nodes.map((node) => ({ + ...node, + selected: node.id === existing.id, + })), existing.id, nextConfig, state.layoutDirection, @@ -444,7 +452,13 @@ export const useRecipeStudioStore = create((set, get) => ({ if (state.executionLocked) { return state; } - const added = buildAddedNodeState(state, "llm", type, position, openDialog); + const added = buildAddedNodeState( + state, + "llm", + type, + position, + openDialog, + ); const context = getAddedNodeContext(added); if (!context) { return added; @@ -495,9 +509,7 @@ export const useRecipeStudioStore = create((set, get) => ({ let { nodes, configs } = context; let edges = state.edges; const unboundModelConfigs = Object.values(configs).filter( - (config) => - config.kind === "model_config" && - !config.provider.trim(), + (config) => config.kind === "model_config" && !config.provider.trim(), ); if (!position && unboundModelConfigs.length > 0) { nodes = placeNodeNear( @@ -605,7 +617,7 @@ export const useRecipeStudioStore = create((set, get) => ({ let { nodes, configs } = context; let edges = state.edges; const unboundLlms = Object.values(configs).filter( - (config) => config.kind === "llm" && !(config.tool_alias?.trim()), + (config) => config.kind === "llm" && !config.tool_alias?.trim(), ); if (!position && unboundLlms.length > 0) { nodes = placeNodeNear( @@ -757,17 +769,15 @@ export const useRecipeStudioStore = create((set, get) => ({ if (cfg.kind !== "model_config" || cfg.provider !== providerName) { continue; } - if (nextIsLocal && !cfg.model.trim()) { - // external -> local: auto fill the placeholder model id so the - // config does not fail "model is required" validation. - configs = { ...configs, [cfgId]: { ...cfg, model: "local" } }; - continue; - } - if (!nextIsLocal && cfg.model === "local") { - // local -> external: clear the placeholder so the user picks a - // real model id for the new external endpoint. - configs = { ...configs, [cfgId]: { ...cfg, model: "" } }; - } + configs = { + ...configs, + [cfgId]: { + ...cfg, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }, + }; } } } diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index 9c720a06d3..b8ed13f70b 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -264,6 +264,8 @@ export type ModelConfig = { kind: "model_config"; name: string; model: string; + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant?: string; provider: string; // biome-ignore lint/style/useNamingConvention: api schema inference_temperature?: string; diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts b/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts index 4fa70b3e3a..059ddf186f 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts @@ -11,7 +11,6 @@ import { isSemanticTargetHandle, normalizeRecipeHandleId, } from "../handles"; -import { isSemanticRelation } from "./relations"; import { isCategoryConfig, isExpressionConfig, @@ -21,6 +20,7 @@ import { VALIDATOR_OXC_CODE_LANGS, VALIDATOR_SQL_CODE_LANGS, } from "../validators/code-lang"; +import { isSemanticRelation } from "./relations"; function buildTemplateWithRef(template: string, ref: string): string { if (template.includes(ref)) { @@ -157,7 +157,10 @@ function isCompetingIncomingEdge( return source.kind === "sampler" && source.sampler_type === "datetime"; } -function isModelSemanticRelation(source: NodeConfig, target: NodeConfig): boolean { +function isModelSemanticRelation( + source: NodeConfig, + target: NodeConfig, +): boolean { return ( (source.kind === "model_provider" && target.kind === "model_config") || (source.kind === "model_config" && target.kind === "llm") || @@ -181,7 +184,9 @@ function canApplyCodeLangToValidator( if (normalized === "python") { return true; } - return VALIDATOR_SQL_CODE_LANGS.includes(normalized as typeof validator.code_lang); + return VALIDATOR_SQL_CODE_LANGS.includes( + normalized as typeof validator.code_lang, + ); } function countHandleUsage( @@ -333,12 +338,8 @@ export function applyRecipeConnection( if (!isValidRecipeConnection(connection, configs)) { return { edges }; } - const initialSource = connection.source - ? configs[connection.source] - : null; - const initialTarget = connection.target - ? configs[connection.target] - : null; + const initialSource = connection.source ? configs[connection.source] : null; + const initialTarget = connection.target ? configs[connection.target] : null; if (!(initialSource && initialTarget)) { return { edges }; } @@ -386,17 +387,36 @@ export function applyRecipeConnection( nextBaseEdges, ); if (source.kind === "model_provider" && target.kind === "model_config") { - // Keep the model_config.model field in sync with provider mode when the - // link is changed via graph drag (the model-config dialog path has its - // own applyProviderChange helper that does the same thing). + // Keep model_config.provider in sync when a graph drag changes the link. + // Local providers now require an explicit selected load id; do not synthesize + // the legacy "local" placeholder. External relinks clear local-only GGUF + // metadata, while legacy placeholders are normalized back to empty. const isSourceLocal = source.is_local === true; - let nextModel = target.model; - if (isSourceLocal && !nextModel.trim()) { - nextModel = "local"; - } else if (!isSourceLocal && nextModel === "local") { - nextModel = ""; - } - const next = { ...target, provider: source.name, model: nextModel }; + const isLegacyLocalPlaceholder = + target.model.trim().toLowerCase() === "local"; + const previousProviderName = target.provider.trim(); + const previousProvider = Object.values(configs).find( + (config) => + config.kind === "model_provider" && + config.name === previousProviderName, + ); + const wasLinkedToLocal = + previousProvider?.kind === "model_provider" && + previousProvider.is_local === true; + const shouldClearModel = + isLegacyLocalPlaceholder || + (isSourceLocal ? !wasLinkedToLocal : wasLinkedToLocal); + const next = { + ...target, + provider: source.name, + ...(shouldClearModel ? { model: "" } : {}), + ...(shouldClearModel || !isSourceLocal + ? { + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + } + : {}), + }; return { edges: nextEdges, configs: { ...configs, [target.id]: next } }; } if (source.kind === "model_config" && target.kind === "llm") { @@ -435,10 +455,9 @@ export function applyRecipeConnection( // biome-ignore lint/style/useNamingConvention: api schema target_columns: [source.name], // biome-ignore lint/style/useNamingConvention: api schema - code_lang: - ( - canUseCodeLangForTarget ? nextCodeLang : target.code_lang - ) as typeof target.code_lang, + code_lang: (canUseCodeLangForTarget + ? nextCodeLang + : target.code_lang) as typeof target.code_lang, }; return { edges: nextEdges, configs: { ...configs, [target.id]: next } }; } diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts index 15ecf39a7b..6b3df846a1 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts @@ -1,15 +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 -import type { - ModelConfig, - ModelProviderConfig, -} from "../../../types"; -import { - isRecord, - readNumberString, - readString, -} from "../helpers"; +import type { ModelConfig, ModelProviderConfig } from "../../../types"; +import { isRecord, readNumberString, readString } from "../helpers"; export function parseModelProvider( provider: Record, @@ -53,6 +46,8 @@ export function parseModelConfig( kind: "model_config", name, model: readString(model.model) ?? "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: readString(model.gguf_variant) ?? undefined, provider: readString(model.provider) ?? "", // biome-ignore lint/style/useNamingConvention: api schema inference_temperature: readNumberString(inference.temperature), diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts index 14e0faa5cc..1575919705 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts @@ -54,55 +54,62 @@ export function buildModelProvider( }; } -export function buildModelConfig( +function assignFiniteNumber( + target: Record, + key: string, + rawValue: string | undefined, + transform: (value: number) => number = (value) => value, +): void { + const trimmed = rawValue?.trim(); + if (!trimmed) { + return; + } + + const parsed = Number(trimmed); + if (Number.isFinite(parsed)) { + target[key] = transform(parsed); + } +} + +function buildInferenceParameters( config: ModelConfig, errors: string[], ): Record { const inference: Record = {}; - const temp = config.inference_temperature?.trim(); - const topP = config.inference_top_p?.trim(); - const maxTokens = config.inference_max_tokens?.trim(); - const timeout = config.inference_timeout?.trim(); + assignFiniteNumber(inference, "temperature", config.inference_temperature); + assignFiniteNumber(inference, "top_p", config.inference_top_p); + assignFiniteNumber(inference, "max_tokens", config.inference_max_tokens); + assignFiniteNumber( + inference, + "timeout", + config.inference_timeout, + Math.trunc, + ); + const extraBody = parseJsonObject( config.inference_extra_body, `Model ${config.name} inference extra_body`, errors, ); - - if (temp) { - const parsed = Number(temp); - if (Number.isFinite(parsed)) { - inference.temperature = parsed; - } - } - if (topP) { - const parsed = Number(topP); - if (Number.isFinite(parsed)) { - // biome-ignore lint/style/useNamingConvention: api schema - inference.top_p = parsed; - } - } - if (maxTokens) { - const parsed = Number(maxTokens); - if (Number.isFinite(parsed)) { - // biome-ignore lint/style/useNamingConvention: api schema - inference.max_tokens = parsed; - } - } - if (timeout) { - const parsed = Number(timeout); - if (Number.isFinite(parsed)) { - inference.timeout = Math.trunc(parsed); - } - } if (extraBody) { - // biome-ignore lint/style/useNamingConvention: api schema inference.extra_body = extraBody; } + return inference; +} + +export function buildModelConfig( + config: ModelConfig, + errors: string[], +): Record { + const inference = buildInferenceParameters(config, errors); + const ggufVariant = config.gguf_variant?.trim(); + return { alias: config.name, model: config.model, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: ggufVariant || undefined, provider: config.provider || undefined, // biome-ignore lint/style/useNamingConvention: api schema inference_parameters: diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts b/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts index 7e72e8d919..a3b3763291 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts @@ -54,7 +54,9 @@ export function validateTimedeltaConfigs( } const reference = config.reference_column_name?.trim() ?? ""; if (!reference) { - errors.push(`Timedelta ${config.name}: reference datetime column required.`); + errors.push( + `Timedelta ${config.name}: reference datetime column required.`, + ); continue; } const parent = nameToConfig.get(reference); @@ -63,7 +65,9 @@ export function validateTimedeltaConfigs( parent.kind !== "sampler" || parent.sampler_type !== "datetime" ) { - errors.push(`Timedelta ${config.name}: reference '${reference}' must be datetime.`); + errors.push( + `Timedelta ${config.name}: reference '${reference}' must be datetime.`, + ); } } } @@ -91,9 +95,18 @@ export function validateModelConfigProviders( const provider = config.provider.trim(); const alias = config.name; const isLocal = localProviderNames.has(provider); - // Local providers do not require a real model id - the loaded Chat - // model is used regardless of what gets sent in the payload. - if (!isLocal && modelAliases.has(alias) && !config.model.trim()) { + const isUsed = modelAliases.has(alias); + const model = config.model.trim(); + const isLegacyLocalPlaceholder = model.toLowerCase() === "local"; + + if (!isLocal && isUsed && isLegacyLocalPlaceholder) { + errors.push(`Model config ${alias}: model is required.`); + continue; + } + if (isLocal && isUsed && !model) { + errors.push(`Model config ${alias}: choose a local model.`); + } + if (!isLocal && isUsed && !model) { errors.push(`Model config ${alias}: model is required.`); } if (provider && !modelProviderNames.has(provider)) { @@ -121,7 +134,9 @@ export function validateUsedProviders( errors.push(`Model provider ${provider.name}: endpoint is required.`); } if (!provider.provider_type.trim()) { - errors.push(`Model provider ${provider.name}: provider_type is required.`); + errors.push( + `Model provider ${provider.name}: provider_type is required.`, + ); } } } @@ -145,7 +160,9 @@ export function validateValidatorConfigs( continue; } if (targetConfig.kind !== "llm" || targetConfig.llm_type !== "code") { - errors.push(`Validator ${config.name}: target '${target}' must be LLM Code.`); + errors.push( + `Validator ${config.name}: target '${target}' must be LLM Code.`, + ); continue; } if ( From b1ef65c07a252e0aabe0986e596b24cf03971721 Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Tue, 26 May 2026 13:17:43 +0200 Subject: [PATCH 18/43] Improve image generation UI (#5784) * Improve image generation UI * Polish generated image edit UI * Tune generated image UI polish * Soften generated image UI * Refine generated image loading surface * Align generated image caption clamp --- .../src/components/assistant-ui/thread.tsx | 18 +-- .../assistant-ui/tool-ui-image-generation.tsx | 108 ++++++++++++++++-- .../src/features/chat/api/chat-adapter.ts | 21 +++- .../src/features/chat/shared-composer.tsx | 21 +++- 4 files changed, 143 insertions(+), 25 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 6a99f30508..da243f37e1 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -75,7 +75,6 @@ import { DownloadIcon, GlobeIcon, HeadphonesIcon, - ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, @@ -89,6 +88,7 @@ import { Copy01Icon, Delete02Icon, Edit03Icon, + Image03Icon, Tick02Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -271,17 +271,17 @@ const GeneratedImageViewportOverlay: FC<{ hideComposer?: boolean }> = ({ className="w-full max-w-[min(100%,46rem)] shrink-0 text-center" title={overlay.title} > -

+

Generated image

{overlay.metadata ? ( -

+

{overlay.metadata}

) : null} {hideComposer ? null : ( -

- Type edits below, then send. +

+ Type edits below, then send

)} @@ -525,13 +525,15 @@ const Composer: FC<{ { : "Enable image generation" } > - + Images ); diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx index e32d23acf6..7dfdd903fe 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx @@ -8,7 +8,7 @@ import { cn } from "@/lib/utils"; import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; import { DownloadIcon, ImageIcon, PencilIcon } from "lucide-react"; import type { CSSProperties, MouseEvent } from "react"; -import { memo, useState } from "react"; +import { memo, useCallback, useEffect, useRef, useState } from "react"; import { useGeneratedImageOverlay } from "./generated-image-overlay-context"; import { Image, downloadImagePart } from "./image"; import { @@ -63,6 +63,8 @@ type GeneratedImagePart = { filename?: string; }; +const CAPTION_COLLAPSED_LINES = 4; + const extensionForMime = (mime: string): string => { switch (mime.toLowerCase()) { case "image/jpeg": @@ -119,7 +121,7 @@ function GeneratedImagePlaceholder({ label }: { label: string }) { return (
220; const imageMetadata = [imageResult?.size, imageResult?.quality, mime] .filter(Boolean) .join(" · "); @@ -174,8 +178,63 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ : null; const [open, setOpen] = useState(true); + const [expandedCaptionPrompt, setExpandedCaptionPrompt] = useState< + string | null + >(null); + const [promptOverflow, setPromptOverflow] = useState<{ + prompt: string; + canExpand: boolean; + } | null>(null); + const captionRef = useRef(null); const isPendingImage = !imagePart && status?.type === "running"; + const promptOverflowMeasured = promptOverflow?.prompt === captionPrompt; + const promptCanExpand = promptOverflowMeasured + ? promptOverflow.canExpand + : false; + const promptExpanded = expandedCaptionPrompt === captionPrompt; + + const updatePromptOverflow = useCallback(() => { + const captionElement = captionRef.current; + if (!captionElement || !captionPrompt) { + return; + } + const computedStyle = window.getComputedStyle(captionElement); + const lineHeight = Number.parseFloat(computedStyle.lineHeight); + const collapsedHeight = + (Number.isFinite(lineHeight) ? lineHeight : 20) * + CAPTION_COLLAPSED_LINES; + const hasOverflow = captionElement.scrollHeight > collapsedHeight + 1; + setPromptOverflow((current) => + current?.prompt === captionPrompt && current.canExpand === hasOverflow + ? current + : { prompt: captionPrompt, canExpand: hasOverflow }, + ); + }, [captionPrompt]); + + useEffect(() => { + const captionElement = captionRef.current; + if (!captionElement || !captionPrompt) { + return; + } + const frame = window.requestAnimationFrame(updatePromptOverflow); + const resizeObserver = + typeof ResizeObserver === "undefined" + ? null + : new ResizeObserver(updatePromptOverflow); + resizeObserver?.observe(captionElement); + window.addEventListener("resize", updatePromptOverflow); + return () => { + window.cancelAnimationFrame(frame); + resizeObserver?.disconnect(); + window.removeEventListener("resize", updatePromptOverflow); + }; + }, [captionPrompt, updatePromptOverflow]); + + const shouldClampPrompt = + (promptOverflowMeasured ? promptCanExpand : promptLikelyNeedsExpansion) && + !promptExpanded; + const runningLabel = "Generating image…"; const completedLabel = formatGeneratedImageLabel(prompt); @@ -231,21 +290,28 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ {imagePart ? (
-
+
+ +
-
+
- {prompt ? ( -
- {prompt} + {captionPrompt ? ( +
+
+ {captionPrompt} +
+ {promptCanExpand ? ( + + ) : null}
) : null}
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 49a5eebd6b..49622c8090 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1433,6 +1433,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Tool call content parts — accumulated and yielded cumulatively. // result is set directly on the tool-call part when tool_end arrives. const toolCallParts: ToolCallMessagePart[] = []; + const orderAssistantContent = ( + textParts: ReturnType, + ) => { + const imageToolParts = toolCallParts.filter( + (part) => part.toolName === "image_generation", + ); + const otherToolParts = toolCallParts.filter( + (part) => part.toolName !== "image_generation", + ); + return [...otherToolParts, ...textParts, ...imageToolParts]; + }; // Anthropic document_citations tool_event payload, converted to // Sources-panel source parts at end-of-stream so the inline [N] // markers have matching entries. @@ -2036,10 +2047,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }; } } - // Yield cumulative state so tool UI updates (tools first, text after) + // Yield cumulative state so tool UI updates. Search/code tools stay + // before the text, while generated images sit after the answer. const textParts = parseAssistantContent(cumulativeText); yield { - content: [...toolCallParts, ...textParts], + content: orderAssistantContent(textParts), metadata: { timing: buildTiming( streamStartTime, @@ -2187,7 +2199,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (parts.length > 0 || toolCallParts.length > 0) { yield { - content: [...toolCallParts, ...parts], + content: orderAssistantContent(parts), metadata: { timing: buildTiming( streamStartTime, @@ -2283,8 +2295,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { yield { content: [ - ...toolCallParts, - ...parseAssistantContent(cumulativeText), + ...orderAssistantContent(parseAssistantContent(cumulativeText)), ...sourceParts, ...documentCitationParts, ], diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 76e77d1288..ba4a897f63 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -21,7 +21,20 @@ import { isTauri } from "@/lib/api-base"; import { isMultimodalResponse } from "./types/api"; import { getImageInputUnavailableReason } from "./utils/image-input-support"; import { useAui } from "@assistant-ui/react"; -import { ArrowUpIcon, DownloadIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; +import { + ArrowUpIcon, + DownloadIcon, + GlobeIcon, + HeadphonesIcon, + LightbulbIcon, + LightbulbOffIcon, + MicIcon, + PlusIcon, + SquareIcon, + XIcon, +} from "lucide-react"; +import { Image03Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { toast } from "@/lib/toast"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers"; @@ -1115,7 +1128,11 @@ export function SharedComposer({ imageToolsEnabled ? "Disable image generation" : "Enable image generation" } > - + Images )} From 41d24227cd2dfef929488854370b53890d21cc4f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 04:31:22 -0700 Subject: [PATCH 19/43] Studio: per-card web_search result + shell_call output fallback (OpenAI) (#5785) * Studio: per-card web_search result + shell_call output fallback (OpenAI) Two empty-output bugs in the OpenAI Responses tool-result rendering that showed up clearly when a single prompt invoked 9 web_search + 4 code_execution + 1 image_generation in one turn. Reproduction shape in the SQLite-stored chat history: - 8 of 9 web_search tool-call records had result == "" (the cards rendered as empty cards in the thread) - 4 of 4 code_execution (shell_call) records were missing the result key entirely (NoneType), so the cards that showed "Ran cat ..." style commands displayed the command line but no output panel at all - image_generation worked, as did the very last web_search of the run Root causes in studio/backend/core/inference/external_provider.py: 1. web_search_call's tool_end emitted result: "" by design, with the intent of overwriting only the LAST call at response.completed with the full citation list (the source-pill extractor on the frontend flatMaps across every web_search result, so a single non-empty result is enough for the trailing source pills). Side effect: every intermediate card renders empty in the thread. Fix: seed each call's own tool_end result with "Searching: " so the per-card text is never empty, then keep the last-call overwrite path so the source-pill extractor still works. Falls back to empty when the model emits an action with no query, so the existing last-call path stays unchanged for that edge. 2. shell_call's tool_start was emitted from response.output_item.done for the call item, but tool_end lived in the separate response.output_item.done handler for shell_call_output. When OpenAI's Responses stream bundles the output array onto the shell_call item's own done event (no separate shell_call_output item), the previous handler emitted tool_start with no following tool_end. The card spun on "running" indefinitely and stored as NoneType in the thread DB. Fix: when the shell_call's done event carries an embedded output list, emit tool_end immediately from that. Track tool_end_emitted on the shell_calls map so a subsequent shell_call_output event (some streams ship both) is skipped instead of double-completing the card. A final flush at response.completed emits tool_end for any orphan shell_call that received neither bundled output nor a separate output event, so cards always finalise. Tests (studio/backend/tests/test_openai_tool_result_fallbacks.py, 6 new): - web_search: three calls, each card's result is its own Searching: query (no empties) - web_search: last call still gets the aggregated citation block when url_citations arrive (pins the overwrite path) - web_search: empty action.query falls back to result == "" (no junk Searching: placeholder) - shell_call: bundled output on done emits a single tool_end with that output as the result text - shell_call: bundled-then-separate output does not double-emit tool_end (subsequent shell_call_output is skipped) - shell_call: orphan call with neither bundled nor separate output is flushed at response.completed so the card finalises 15/15 tests green when combined with the existing 9 in test_openai_code_execution.py. Pre-commit + ruff format clean. Scope: OpenAI Responses-API code path only. The Anthropic native Messages-API path (_stream_anthropic) is untouched, as is the local llama-server path. Local-model behaviour cannot regress because the edited handlers only fire inside the OpenAI cloud branch. * Studio: per-model external max_tokens cap + clamp on model switch Two related external-provider issues that surfaced from the same investigation as the per-card web_search / shell_call result bugs in the previous commit: A. Slider cap was a one-size-fits-all 32768 for every external model. provider-capabilities.ts kept a single EXTERNAL_MAX_OUTPUT_TOKENS constant (32k), well below what most providers actually accept. The docstring even called out the right per-provider numbers (Anthropic Opus 128k, GPT-5.x ~128k, Gemini 2.5 ~65k, DeepSeek 8k) but the code picked the lowest as a conservative floor. Effect: long generations from gpt-5.5 / claude-opus-4-7 silently truncated at 32k even though the API would have served up to 128k. Fix: introduce getExternalMaxOutputTokens(providerType, modelId) returning the documented per-model cap. Patterns are checked longest-first so e.g. gpt-5.5-pro matches before gpt-5.5. Unknown provider/model combinations fall back to the existing 32k floor so no surprise increases for ids we don't know about. Per-model caps from the official docs: - OpenAI gpt-5.5 / gpt-5.5-pro: 128000 - OpenAI gpt-5.4 / gpt-5.4-pro: 65536 - OpenAI gpt-5.3: 16384 - Anthropic claude-opus-4-7: 128000 - Anthropic claude-opus-4-6 / sonnet-4-6 / opus-4-5 / sonnet-4-5 / haiku-4-5: 64000 - Gemini 3.x family: 65535 - DeepSeek: 8192 - OpenRouter: strip provider/ prefix from the id and re-resolve The slider in chat-settings-sheet.tsx and the send-time clamp in chat-adapter.ts both call the new function so the slider's max= matches what the wire layer will accept. B. Slider value lied after switching from a local model to external. When Studio auto-loads the helper Gemma-4-E2B-it on first chat, chat-adapter sets params.maxTokens to Gemma's context_length (262144 for Gemma 4). Switching the model picker to gpt-5.5 then flips the slider's max prop to the external cap, but the stored params.maxTokens is never reset. The numeric value next to the slider would render 262144 against a track that ended at the external cap. The send-time clamp brought the outbound max_tokens back down to the cap, so the API call was safe, but the displayed number had no relationship to what was actually being sent. Fix: chat-runtime-store.setCheckpoint now clamps params.maxTokens to getExternalMaxOutputTokens(...) on transitions into an external model. Looks up the provider via useExternalProvidersStore so we can derive providerType from the parsed external model id. No-op when the stored maxTokens is already at or below the new cap, so user-tuned values within range survive the switch. Scope: pure frontend changes scoped to external-provider code paths. Local model behaviour is untouched -- the ggufContextLength branch of the slider's max= is unchanged, and setCheckpoint only mutates maxTokens when isExternalModelId(modelId) is true. The send-time clamp continues to be the safety net for any in-flight request that crosses a model switch before the store-level clamp has applied. Typecheck (tsc -b) clean; bun run build succeeds (2.13s). Co-changes with the previous commit (7fe1adbf, per-card web_search + shell_call output fallback) form a single PR: every empty-output and silent-truncation issue surfaced from the same animal-popularity prompt reproduction is now addressed in one branch. * Studio: correct external max_tokens caps for Gemini and DeepSeek Per-doc corrections to the per-model cap table added in 95da8d52: - Gemini 3.x family: 65535 -> 65536, per https://ai.google.dev/gemini-api/docs/models/gemini-3.1-pro-preview (the published max_output_tokens is exactly 64K = 65536). The earlier 65535 was an off-by-one rough cap. - DeepSeek (deepseek-chat / deepseek-reasoner aliases): 8192 -> 384000, per https://api-docs.deepseek.com/quick_start/pricing. DeepSeek V4 Flash / Pro both list MAX OUTPUT = 384K; the chat / reasoner ids are deprecated aliases for V4 Flash non-thinking / thinking modes. The 8192 value was carried over from V3 and silently truncated V4 traffic at 2% of its actual ceiling. Affects only the slider max and the send-time clamp for these provider types. Other providers' caps unchanged. tsc -b clean. * Studio: also flush orphan shell_calls on response.incomplete Addresses gemini-code-assist[bot] high-priority inline review on PR 5785: the orphan-shell_call final flush added in 7fe1adbf landed only in the response.completed branch. Truncated OpenAI Responses streams emit response.incomplete instead (for example when the request hits max_output_tokens), which left in-flight shell_call cards spinning indefinitely in the UI. Mirror the same flush block in the response.incomplete handler so the truncated-stream path finalizes every pending tool card. The tool_end_emitted guard keeps the path idempotent: if a shell_call already completed via bundled output on its done event, the incomplete flush is a no-op for it. Two new tests in test_openai_tool_result_fallbacks.py: - test_shell_call_flushed_on_response_incomplete_truncation pins the bug repro: an in-flight shell_call followed by response.incomplete must emit tool_end so the card finalizes. - test_shell_call_incomplete_does_not_double_emit pins idempotency: a shell_call that completed via bundled output and is then followed by response.incomplete emits exactly one tool_end with the bundled result text. 17/17 tests green (8 fallback tests + 9 existing code-execution). Pre- commit + ruff format clean. * Studio: trim verbose comments across PR 5785 edits Compress the in-code commentary added across this branch to one or two lines per block; the verbose prose was easier as a PR description than as inline noise. No behavioural changes: 17/17 tests still green, tsc -b still clean. --- .../core/inference/external_provider.py | 85 +++- .../test_openai_tool_result_fallbacks.py | 372 ++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 14 +- .../src/features/chat/chat-settings-sheet.tsx | 6 +- .../features/chat/provider-capabilities.ts | 95 ++++- .../chat/stores/chat-runtime-store.ts | 24 +- 6 files changed, 564 insertions(+), 32 deletions(-) create mode 100644 studio/backend/tests/test_openai_tool_result_fallbacks.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index f34ef5fd62..8f34bb23fc 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -3873,14 +3873,16 @@ class ExternalProviderClient: ), } ) + # Per-card text; last call gets overwritten + # with citations at response.completed. + per_call_result = ( + f"Searching: {query}" if query else "" + ) yield _emit_tool_event( { "type": "tool_end", "tool_call_id": item_id, - # Empty result — the last call gets - # overwritten with citations at - # response.completed. - "result": "", + "result": per_call_result, } ) elif item.get("type") == "shell_call": @@ -3908,7 +3910,11 @@ class ExternalProviderClient: ) shell_calls.setdefault( item_id, - {"commands": [], "output": None}, + { + "commands": [], + "output": None, + "tool_end_emitted": False, + }, ) shell_calls[item_id]["commands"] = ( list(commands) @@ -3926,6 +3932,24 @@ class ExternalProviderClient: }, } ) + # Fallback: output may be bundled on the + # shell_call done event itself. + embedded_output = item.get("output") + if ( + isinstance(embedded_output, list) + and embedded_output + ): + shell_calls[item_id]["output"] = embedded_output + shell_calls[item_id]["tool_end_emitted"] = True + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": item_id, + "result": _format_shell_output( + embedded_output + ), + } + ) elif item.get("type") == "shell_call_output": # `call_id` links back to the shell_call's # `id`, which is what we used as the @@ -3936,8 +3960,15 @@ class ExternalProviderClient: item.get("call_id") or item.get("id") or "" ) output = item.get("output") or [] + # Skip if bundled-output path already + # finalised this card. + if shell_calls.get(call_id, {}).get( + "tool_end_emitted" + ): + continue if call_id in shell_calls: shell_calls[call_id]["output"] = output + shell_calls[call_id]["tool_end_emitted"] = True result_text = _format_shell_output(output) yield _emit_tool_event( { @@ -4093,15 +4124,10 @@ class ExternalProviderClient: } ) container_id_emitted = True - # Apply the aggregated citation list onto the - # *last* web_search call by overwriting its - # tool_end result. The frontend's - # parseSourcesFromResult flatMaps every - # web_search tool-call result, so a single - # non-empty result is enough to surface the - # whole source-pill set at the message tail — - # no need to fan out across every card (which - # would just duplicate the same pills). + # Overwrite the last web_search call with the + # citation list; the source-pill extractor + # flatMaps across cards. Earlier cards keep + # their per-call "Searching:" text. if web_search_calls and all_url_citations: last_id = list(web_search_calls.keys())[-1] blocks: list[str] = [] @@ -4119,6 +4145,21 @@ class ExternalProviderClient: "result": "\n---\n".join(blocks), } ) + # Final flush: finalise any orphan shell_call + # so the card stops spinning. + for sc_id, sc_state in shell_calls.items(): + if sc_state.get("tool_end_emitted"): + continue + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": sc_id, + "result": _format_shell_output( + sc_state.get("output") or [] + ), + } + ) + sc_state["tool_end_emitted"] = True chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -4197,6 +4238,22 @@ class ExternalProviderClient: "result": "\n---\n".join(blocks), } ) + # Mirror the response.completed flush so + # truncated streams also finalise orphan + # shell_calls. + for sc_id, sc_state in shell_calls.items(): + if sc_state.get("tool_end_emitted"): + continue + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": sc_id, + "result": _format_shell_output( + sc_state.get("output") or [] + ), + } + ) + sc_state["tool_end_emitted"] = True chunk = { "id": completion_id, "object": "chat.completion.chunk", diff --git a/studio/backend/tests/test_openai_tool_result_fallbacks.py b/studio/backend/tests/test_openai_tool_result_fallbacks.py new file mode 100644 index 0000000000..7c033bc348 --- /dev/null +++ b/studio/backend/tests/test_openai_tool_result_fallbacks.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for OpenAI Responses tool-result rendering. + +Covers two bug classes: empty web_search cards (per-card result seeded +with "Searching: ") and orphan shell_call cards (bundled-output +fallback + final flush at response.completed / response.incomplete). +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +async def _collect(agen): + out = [] + async for line in agen: + out.append(line) + return out + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client(base_url: str = "https://api.openai.com/v1") -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "openai", + base_url = base_url, + api_key = "sk-test", + ) + + +def _openai_sse(events: list[dict]) -> bytes: + chunks: list[str] = [] + for event in events: + chunks.append(f"event: {event['type']}") + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _tool_events(lines: list[str]) -> list[dict]: + out: list[dict] = [] + for line in lines: + if not line.startswith("data:"): + continue + raw = line[len("data:") :].strip() + if not raw or raw == "[DONE]": + continue + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict) and "_toolEvent" in parsed: + out.append(parsed["_toolEvent"]) + return out + + +def _drive_stream(sse_events, enabled_tools, monkeypatch): + def handler(request): + return httpx.Response( + 200, + content = _openai_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "x"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = enabled_tools, + ) + ) + + return _drive(run()) + + +# ── web_search per-card result ───────────────────────────────────────── + + +def test_web_search_each_call_carries_its_own_query_as_result(monkeypatch): + """Each card carries its own `Searching: ` text; no empties.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_1", + "action": {"query": "popular animals 2026"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_2", + "action": {"query": "most loved animals poll"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_3", + "action": {"query": "tiger ranking"}, + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["web_search"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + by_id = {e["tool_call_id"]: e for e in ends} + assert by_id["ws_1"]["result"] == "Searching: popular animals 2026" + assert by_id["ws_2"]["result"] == "Searching: most loved animals poll" + assert by_id["ws_3"]["result"] == "Searching: tiger ranking" + + +def test_web_search_last_call_overwritten_with_citations(monkeypatch): + """Last call still gets the aggregated citation list; earlier calls + keep their per-call `Searching:` text.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_1", + "action": {"query": "first query"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_2", + "action": {"query": "second query"}, + }, + }, + { + "type": "response.output_text.annotation.added", + "annotation": { + "type": "url_citation", + "url": "https://example.com/a", + "title": "Example A", + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["web_search"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + by_id: dict = {} + # Keep the LAST tool_end per id (the citation overwrite for ws_2). + for e in ends: + by_id[e["tool_call_id"]] = e + # First call keeps its own query. + assert by_id["ws_1"]["result"] == "Searching: first query" + # Last call gets overwritten with the citation block. + assert "Title: Example A" in by_id["ws_2"]["result"] + assert "URL: https://example.com/a" in by_id["ws_2"]["result"] + + +def test_web_search_empty_query_falls_back_to_empty_result(monkeypatch): + """No query -> empty result (no `Searching:` placeholder).""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_only", + "action": {}, + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["web_search"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert len(ends) == 1 + assert ends[0]["result"] == "" + + +# ── shell_call output fallbacks ──────────────────────────────────────── + + +def test_shell_call_emits_tool_end_when_output_bundled_on_done(monkeypatch): + """Output bundled on the shell_call done event emits tool_end.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_bundled", + "action": {"commands": ["echo hi"]}, + "output": [ + { + "stdout": "hi\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + starts = [e for e in events if e["type"] == "tool_start"] + ends = [e for e in events if e["type"] == "tool_end"] + assert len(starts) == 1 + assert starts[0]["tool_call_id"] == "scall_bundled" + assert len(ends) == 1 + assert ends[0]["tool_call_id"] == "scall_bundled" + assert "hi" in ends[0]["result"] + + +def test_shell_call_bundled_then_separate_output_does_not_double_emit(monkeypatch): + """Separate shell_call_output after bundled-output is a no-op.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_both", + "action": {"commands": ["echo bundle"]}, + "output": [ + { + "stdout": "bundle\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call_output", + "id": "scout_both", + "call_id": "scall_both", + "output": [ + { + "stdout": "should not double-emit\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert len(ends) == 1 + assert ends[0]["tool_call_id"] == "scall_both" + assert "bundle" in ends[0]["result"] + assert "should not double-emit" not in ends[0]["result"] + + +def test_shell_call_final_flush_on_completed_when_no_output_event(monkeypatch): + """Orphan shell_call finalises via the response.completed flush.""" + sse_events = [ + { + "type": "response.output_item.added", + "item": { + "type": "shell_call", + "id": "scall_orphan", + "action": {"commands": ["true"]}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_orphan", + "action": {"commands": ["true"]}, + "status": "completed", + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert any(e["tool_call_id"] == "scall_orphan" for e in ends) + + +def test_shell_call_flushed_on_response_incomplete_truncation(monkeypatch): + """Truncated streams (response.incomplete) also flush orphan calls.""" + sse_events = [ + { + "type": "response.output_item.added", + "item": { + "type": "shell_call", + "id": "scall_truncated", + "action": {"commands": ["long_running"]}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_truncated", + "action": {"commands": ["long_running"]}, + "status": "in_progress", + }, + }, + { + "type": "response.incomplete", + "response": { + "incomplete_details": {"reason": "max_output_tokens"}, + }, + }, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert any(e["tool_call_id"] == "scall_truncated" for e in ends) + + +def test_shell_call_incomplete_does_not_double_emit(monkeypatch): + """response.incomplete is idempotent against already-finalised calls.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_done", + "action": {"commands": ["echo done"]}, + "output": [ + { + "stdout": "done\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + { + "type": "response.incomplete", + "response": { + "incomplete_details": {"reason": "max_output_tokens"}, + }, + }, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert len(ends) == 1 + assert ends[0]["tool_call_id"] == "scall_done" + assert "done" in ends[0]["result"] diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 49622c8090..ce59429762 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -21,6 +21,7 @@ import { pickFriendlyContainerName } from "../lib/friendly-names"; import { EXTERNAL_MAX_OUTPUT_TOKENS, clampReasoningEffortToLevels, + getExternalMaxOutputTokens, getExternalMinOutputTokens, getExternalReasoningCapabilities, getProviderCapabilities, @@ -1703,18 +1704,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(externalCapabilities?.topP !== false ? { top_p: params.topP } : {}), - // Clamp to the cross-provider output cap so a maxTokens value - // carried over from a local-model session does not blow past - // provider limits (e.g. Claude Opus 400s on >128k). Also - // floor to the provider's documented minimum — Kimi's - // thinking models need >=16k or the response truncates - // before the answer fits alongside reasoning_content. + // Floor at the provider's documented min (Kimi thinking + // needs >=16k); clamp at the per-model max. max_tokens: Math.min( Math.max( params.maxTokens, getExternalMinOutputTokens(externalProvider?.providerType), ), - EXTERNAL_MAX_OUTPUT_TOKENS, + getExternalMaxOutputTokens( + externalProvider?.providerType, + externalSelection?.modelId, + ), ), // Only forward sampling knobs the provider actually accepts; the // backend's external-provider proxy is param-permissive and would diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index ac1ef8a24c..3714fb128a 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -85,6 +85,7 @@ import { import { EXTERNAL_MAX_OUTPUT_TOKENS, type ProviderCapabilities, + getExternalMaxOutputTokens, getExternalMinOutputTokens, providerSupportsBuiltinCodeExecution, providerSupportsFastMode, @@ -1309,7 +1310,10 @@ export function ChatSettingsPanel({ } max={ isExternalModel - ? EXTERNAL_MAX_OUTPUT_TOKENS + ? getExternalMaxOutputTokens( + externalProviderType, + externalSelection?.modelId, + ) : isGguf && ggufContextLength ? ggufContextLength : 32768 diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index ef805305be..5adc01ea2b 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -71,18 +71,95 @@ export function clampReasoningEffortToLevels( } /** - * Output-token cap for any external provider request. Picked to stay below the - * tightest declared limit across the providers we ship (Anthropic Claude Opus - * tops out at 128k, GPT-5.x ~128k, Gemini 2.5 ~65k, DeepSeek 8k) while staying - * well above what a typical chat reply needs. The local-model path is not - * subject to this — local backends honour whatever the loaded context allows. - * - * If a user's stored maxTokens (e.g. carried over from a prior local-model - * session with a 128k+ context) exceeds this, chat-adapter clamps the - * outbound request so the provider does not 400 on it. + * Fallback cap for unknown providers / models. Prefer + * `getExternalMaxOutputTokens(providerType, modelId)` for the real cap. */ export const EXTERNAL_MAX_OUTPUT_TOKENS = 32768; +/** + * Per-model max-output caps from each provider's docs: + * OpenAI: developers.openai.com/api/docs/models/gpt-5.5 + * Anthropic: platform.claude.com/docs/en/about-claude/models + * Gemini: ai.google.dev/gemini-api/docs/models/gemini-3.1-pro-preview + * DeepSeek: api-docs.deepseek.com/quick_start/pricing (V4 family) + * Local-model path is unaffected. + */ +const EXTERNAL_MAX_OUTPUT_TOKENS_BY_MODEL: Array<{ + providerType: string; + prefixes: readonly string[]; + cap: number; +}> = [ + // OpenAI + { providerType: "openai", prefixes: ["gpt-5.5-pro", "gpt-5.5"], cap: 128000 }, + { providerType: "openai", prefixes: ["gpt-5.4-pro", "gpt-5.4"], cap: 65536 }, + { providerType: "openai", prefixes: ["gpt-5.3"], cap: 16384 }, + // Anthropic + { + providerType: "anthropic", + prefixes: ["claude-opus-4-7"], + cap: 128000, + }, + { + providerType: "anthropic", + prefixes: [ + "claude-opus-4-6", + "claude-sonnet-4-6", + "claude-opus-4-5", + "claude-sonnet-4-5", + "claude-haiku-4-5", + ], + cap: 64000, + }, + // Gemini + { + providerType: "gemini", + prefixes: ["gemini-3", "gemini-pro", "gemini-flash"], + cap: 65536, + }, + // DeepSeek (V4: deepseek-chat / deepseek-reasoner alias V4-flash). + { providerType: "deepseek", prefixes: ["deepseek"], cap: 384000 }, +]; + +/** + * Documented per-model output cap; unknown ids fall back to + * `EXTERNAL_MAX_OUTPUT_TOKENS` (32k). OpenRouter ids are + * `provider/model`; the prefix is stripped before matching. + */ +export function getExternalMaxOutputTokens( + providerType: string | null | undefined, + modelId: string | null | undefined, +): number { + if (!providerType || !modelId) return EXTERNAL_MAX_OUTPUT_TOKENS; + const normalized = modelId.trim().toLowerCase(); + if (!normalized) return EXTERNAL_MAX_OUTPUT_TOKENS; + const stripped = + providerType === "openrouter" && normalized.includes("/") + ? normalized.split("/").slice(-1)[0] + : normalized; + const effectiveProvider = + providerType === "openrouter" + ? _inferProviderFromOpenrouterId(normalized) ?? providerType + : providerType; + for (const entry of EXTERNAL_MAX_OUTPUT_TOKENS_BY_MODEL) { + if (entry.providerType !== effectiveProvider) continue; + if (entry.prefixes.some((prefix) => stripped.startsWith(prefix))) { + return entry.cap; + } + } + return EXTERNAL_MAX_OUTPUT_TOKENS; +} + +function _inferProviderFromOpenrouterId( + normalizedId: string, +): string | null { + // Map OpenRouter `provider/model` prefix to our internal providerType. + if (normalizedId.startsWith("openai/")) return "openai"; + if (normalizedId.startsWith("anthropic/")) return "anthropic"; + if (normalizedId.startsWith("google/")) return "gemini"; + if (normalizedId.startsWith("deepseek/")) return "deepseek"; + return null; +} + /** * Whether the external provider offers a built-in web-search tool that the * model invokes server-side. When `true`, the chat composer's Search button 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 c78f02a474..a71e4127a2 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -14,7 +14,9 @@ import { DEFAULT_INFERENCE_PARAMS, type InferenceParams, } from "../types/runtime"; -import { isExternalModelId } from "../external-providers"; +import { isExternalModelId, parseExternalModelId } from "../external-providers"; +import { getExternalMaxOutputTokens } from "../provider-capabilities"; +import { useExternalProvidersStore } from "./external-providers-store"; import { loadChatSettingsWithLegacyImport, savePersistedChatSettingsPatch, @@ -747,10 +749,30 @@ export const useChatRuntimeStore = create((set, get) => ({ // external-provider render gate would otherwise show old counters // until the next completion overwrites them. const checkpointChanged = state.params.checkpoint !== modelId; + // Clamp maxTokens to the new model's cap on switch into an + // external model so a value carried over from a prior local + // session does not render above the slider's max. + let nextMaxTokens = state.params.maxTokens; + if (checkpointChanged && isExternalModelId(modelId)) { + const parsed = parseExternalModelId(modelId); + const provider = parsed + ? useExternalProvidersStore + .getState() + .providers.find((p) => p.id === parsed.providerId) + : null; + const cap = getExternalMaxOutputTokens( + provider?.providerType, + parsed?.modelId, + ); + if (nextMaxTokens > cap) { + nextMaxTokens = cap; + } + } return { params: { ...state.params, checkpoint: modelId, + maxTokens: nextMaxTokens, }, activeGgufVariant: ggufVariant ?? null, ...(checkpointChanged ? { contextUsage: null } : {}), From fb65fed3b0f0c4a91ca88043da19b95424ba26c6 Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Tue, 26 May 2026 13:37:18 +0200 Subject: [PATCH 20/43] Keep generated image loading dots animated (#5786) --- studio/frontend/src/index.css | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 8a1fe13678..a3a1e8dad0 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1244,11 +1244,11 @@ * end state. Hover colour changes become instant rather than fading, which is * the documented WCAG outcome (motion is "minimised, not removed"). * - * .animate-spin is the exception: loading spinners are essential progress + * .animate-spin and generated image loading dots are the exceptions: loading * indicators across Studio (tool execution loaders, sonner toasts, Tauri - * startup / update screens, the primitive). Freezing them - * removes the only visual signal that work is in flight, so they keep - * animating but at a slower, less aggressive 1.5s cadence. + * startup / update screens, the primitive, and image generation + * cards). Freezing them removes the only visual signal that work is in flight, + * so they keep animating. */ @media (prefers-reduced-motion: reduce) { *, @@ -1264,4 +1264,9 @@ animation-duration: 1.5s !important; animation-iteration-count: infinite !important; } + + .generated-image-loading-dot { + animation-duration: 1850ms !important; + animation-iteration-count: infinite !important; + } } From c49cc6daf58623032ea552b05cc77bdc8d86573c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 05:29:42 -0700 Subject: [PATCH 21/43] Studio: auto-recover when shadowed 'unsloth' on PATH hides the frontend dist (#5782) * Studio: auto-recover when shadowed 'unsloth' on PATH hides the frontend dist The CLI launcher derives `_PACKAGE_ROOT` from where `unsloth_cli` imports from, and `studio/backend/run.py` derives its default `frontend_path` from `Path(__file__).resolve().parent.parent / "frontend" / "dist"`. When another `unsloth` (a separate venv with `pip install unsloth`, a system install, an older venv earlier on PATH) wins `which unsloth`, both resolve into a site-packages tree that ships frontend source files but no vite-built `dist/`. The backend warned `[WARNING] Frontend not found at ...` and then happily served 200 on every `/api/*` route while returning `{"detail":"Not Found"}` on `/`. The 404 was silent to users -- the process was healthy, the log line scrolled by, and the only symptom was a blank browser tab. This is a real situation: many devboxes carry a workspace venv with `unsloth` installed years before the user runs `curl|sh` to install Studio. The installer-managed binary at `~/.local/bin/unsloth` exists but loses to the older venv on PATH order. Three layers of fix, additive: Layer C -- runtime auto-discovery (unsloth_cli + run.py) The CLI now resolves `--frontend` explicitly before spawning `run.py`, probing in order: package-local default, installer venv site-packages (`$STUDIO_HOME/unsloth_studio/lib/python*/site-packages/...` and the Windows `Lib/site-packages/...` equivalent), and editable-install source roots read from `__editable___*_finder.py` MAPPING dicts in the installer venv. `run.py` does the same probe as a backstop for direct `python run.py` invocations. Layer E -- loud structured error The silent `[WARNING]` is replaced with a `SystemExit` that names every candidate path tried and lists the four one-line fixes (run the absolute path, pass `--frontend`, pass `--api-only`, reinstall). Suppressed only in `--api-only` mode where no UI is served by design. Layer F -- installer self-check (install.sh + install.ps1) At the tail of install, both installers compare `command -v unsloth` (POSIX) / `Get-Command unsloth` (PowerShell) against the just-installed binary. If a different path wins, a yellow `warning` block names the shadowing binary and prints the alias / absolute-path / PATH-reorder fixes. install.sh uses the venv Python for path canonicalization so it also works on macOS (BSD `readlink` has no `-f`). Cross-platform notes: - Glob patterns probe both `lib/python*/site-packages` (POSIX) and `Lib/site-packages` (Windows). - Canonical-binary path branches on `sys.platform == "win32"` to pick `unsloth.exe` over `unsloth`. - install.sh fixed for macOS; install.ps1 is the Windows analog. Tests: `studio/backend/tests/test_frontend_resolution.py` covers five cases via AST-load of the helpers (no uvicorn / FastAPI import needed, matching `test_host_defaults.py`'s style): 1. Resolver returns None when nothing exists anywhere. 2. Resolver picks the first existing candidate when the default works. 3. Fallback to `$UNSLOTH_STUDIO_HOME` site-packages dist when the default is missing. 4. Fallback to an editable-install source root via MAPPING parsing. 5. Resolver tolerates a non-existent `$UNSLOTH_STUDIO_HOME`. All 5 new + 2 existing host-default tests pass. * Studio: address review feedback on PR 5782 (Windows hardlink, Win path hint, broader tests) Four parallel platform reviews (Windows, Linux, macOS, general) on the initial commit surfaced a small batch of correctness items, all addressed here: Windows install.ps1 (medium severity, false positive on every install): The user-facing shim at $StudioHome\bin\unsloth.exe is a hardlink to $VenvDir\Scripts\unsloth.exe (created at line 1582). Resolve-Path does not de-duplicate hardlinks, so the previous string compare always saw the two paths as different and the new "another 'unsloth' wins on PATH" warning would fire on every fresh Windows install. Switched to content-hash equality via Get-FileHash, which collapses hardlinks, symlinks, and identical copies to a single identity. Also restricted the probe to Get-Command -CommandType Application so PowerShell aliases / functions / scripts named "unsloth" don't false-trigger. Windows run.py SystemExit hint (medium severity, defeats the recovery UX): The structured error printed Path(STUDIO_HOME)/"unsloth_studio"/"bin"/ "unsloth.exe" on every platform, but on Windows the installer places the shim at $STUDIO_HOME/bin/unsloth.exe (no unsloth_studio segment) and the venv binary at $STUDIO_HOME/unsloth_studio/Scripts/unsloth.exe. The hint pointed at a non-existent path on Windows. Branch on sys.platform == "win32" to emit the real shim location; Linux / macOS keep the unsloth_ studio/bin/unsloth layout. MAPPING regex robustness (low): [^\n]* silently failed if a future setuptools / black reformat wrapped the MAPPING dict across multiple lines. Tightened to [^}]* + re.DOTALL, which still rejects nested dicts (setuptools never emits those for editable installs) but tolerates either single- or multi-line literals. install.sh broken-venv edge case (low, macOS reviewer): Previously _canon fell back to echoing the raw input when the venv python failed, which would make two symlinked-but-identical paths look different and false-trigger the warning. Now _canon returns empty on failure and the caller skips the whole comparison if either side is unresolvable. argparse default + log readability (nits): run.py's argparse --frontend default now reuses the module-level _DEFAULT_FRONTEND_PATH constant so it stays in lockstep with run_server's default. The [OK] log message resolves the chosen path so support output is always absolute. Tests grow from 5 to 8 in studio/backend/tests/test_frontend_resolution. py (10/10 with the existing host-default tests): - Windows-layout fallback: Lib/site-packages with capital L. - Multi-line MAPPING dict: locks in the [^}]* + re.DOTALL behaviour. - SystemExit message contract: every actionable fix string and the attempted-paths list must appear; pins the user-facing recovery message so a future refactor doesn't drop a bullet. End-to-end re-verified on this box: shadowing workspace_22/bin/unsloth still serves 200 on / through the editable-finder fallback, with the follow-up resolve-then-log change yielding [OK] Frontend loaded from /mnt/disks/unslothai/ubuntu/unsloth/studio/frontend/dist. Out of scope (called out by reviewers but deferred): - _resolve_frontend_path candidate ordering still tries _PACKAGE_ROOT first. For the rare case where a shadowing install carries an older built dist, this serves the stale UI instead of the fresh one. Fix is non-trivial (the --local workflow intentionally wants _PACKAGE_ROOT to win when the cloned repo is the source of truth), so leaving it for a follow-up. - studio/backend/colab.py still bails out on missing frontend instead of routing through the new resolver. Pre-existing behaviour, separate PR. - _resolve_frontend_path is duplicated across run.py and unsloth_cli/ commands/studio.py. Minor maintenance concern; consolidation is natural in a later refactor. * Studio: guard ast.literal_eval result with isinstance(dict) Addresses gemini-code-assist[bot] high-priority inline review on PR 5782 flagging that `mapping.get('studio')` could raise AttributeError if the MAPPING regex matched a brace-delimited literal that ast.literal_eval parsed as a non-dict (set, list, None). The regex `\{[^}]*\}` happily matches `{1, 2, 3}` and literal_eval returns a set; the previous code then crashed on .get(). Setuptools's editable-install template only emits dict literals so this is defensive rather than a live bug, but the guard is one line per call site and prevents a future template change from taking out backend startup or CLI invocation. Both call sites (studio/backend/run.py:558 and unsloth_cli/commands/studio.py:234) now bail out on the finder file when isinstance(mapping, dict) is False; the resolver keeps probing the remaining finders, so a malformed entry in one finder cannot poison the discovery of a good one elsewhere. Adds test_resolver_does_not_crash_on_non_dict_mapping_literal to test_frontend_resolution.py, which writes one bad finder (MAPPING is a set literal) alongside one good finder (MAPPING is a real dict) and asserts the resolver returns the good finder's dist path. Without the guard this test crashes with AttributeError; with the guard it passes. 11/11 tests green. --- install.ps1 | 27 ++ install.sh | 32 +++ studio/backend/run.py | 133 +++++++++- .../backend/tests/test_frontend_resolution.py | 248 ++++++++++++++++++ unsloth_cli/commands/studio.py | 83 +++++- 5 files changed, 514 insertions(+), 9 deletions(-) create mode 100644 studio/backend/tests/test_frontend_resolution.py diff --git a/install.ps1 b/install.ps1 index 3911236d87..52766370d1 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1626,6 +1626,33 @@ shell.Run cmd, 0, False # New-StudioShortcuts gates the .lnk shortcuts on env-mode internally. New-StudioShortcuts -UnslothExePath $UnslothExe + # Warn if another 'unsloth' wins on PATH (different venv, system pip). + # Mirrors install.sh; absolute path is still the most reliable launch. + # Uses content-hash equality (Get-FileHash) so hardlinks, symlinks, and + # identical copies of the installer's shim don't false-trigger. CommandType + # Application restricts the probe to real executables (skips aliases, + # functions, scripts). + try { + $_pathCmd = Get-Command unsloth -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1 + if ($_pathCmd) { + $_pathExe = $_pathCmd.Source + $_installedHash = (Get-FileHash -LiteralPath $UnslothExe -Algorithm SHA256 -ErrorAction SilentlyContinue).Hash + $_pathHash = (Get-FileHash -LiteralPath $_pathExe -Algorithm SHA256 -ErrorAction SilentlyContinue).Hash + if ($_installedHash -and $_pathHash -and ($_installedHash -ne $_pathHash)) { + Write-Host "" + step "warning" "another 'unsloth' wins on PATH:" "Yellow" + substep $_pathExe + substep "this installer's binary is at:" + substep $UnslothExe + substep "to use this install, call the absolute path above," + substep "or put its dir earlier on PATH." + Write-Host "" + } + } + } catch { + # Diagnostic only; never block install on a probe failure. + } + # In interactive terminals, ask the user before starting Studio. # In non-interactive environments (CI, Docker) just print instructions. $IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) diff --git a/install.sh b/install.sh index cc92fd52c2..d12abe298f 100755 --- a/install.sh +++ b/install.sh @@ -2263,6 +2263,38 @@ if [ "$TAURI_MODE" = true ]; then exit 0 fi +# Warn if another 'unsloth' wins on PATH (different venv, system pip, etc). +# Users typing `unsloth studio` later would hit that binary instead of the +# one just installed; the runtime now falls back via UNSLOTH_STUDIO_HOME +# but the absolute path is still the most reliable launch. +# Uses the venv python (just created above) for path canonicalization so +# this works on macOS (BSD readlink has no -f) as well as Linux/WSL. +_installed_bin="$VENV_DIR/bin/unsloth" +_path_unsloth=$(command -v unsloth 2>/dev/null || true) +if [ -n "$_path_unsloth" ] && [ -x "$VENV_DIR/bin/python" ]; then + # Canonicalize via the venv python (BSD readlink lacks -f on macOS). + # If either side fails to resolve, skip the check entirely rather than + # comparing raw paths (which would false-trigger on symlink targets). + _canon() { + "$VENV_DIR/bin/python" -c \ + 'import os, sys; print(os.path.realpath(sys.argv[1]))' \ + "$1" 2>/dev/null + } + _installed_real=$(_canon "$_installed_bin") + _path_real=$(_canon "$_path_unsloth") + if [ -n "$_installed_real" ] && [ -n "$_path_real" ] \ + && [ "$_installed_real" != "$_path_real" ]; then + echo "" + step "warning" "another 'unsloth' wins on PATH:" "$C_WARN" + substep "$_path_unsloth" + substep "this installer's binary is at:" + substep "$_installed_bin" + substep "to use this install, run the absolute path above," + substep "alias unsloth, or put its dir earlier on PATH." + echo "" + fi +fi + echo "" printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" printf " ${C_DIM}%s${C_RST}\n" "$RULE" diff --git a/studio/backend/run.py b/studio/backend/run.py index d5ccc49022..3bde8abd3c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -9,6 +9,7 @@ Works independently and can be moved to any directory. import os import sys from pathlib import Path +from typing import Optional # Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked) os.environ["PYTHONWARNINGS"] = "ignore" @@ -512,10 +513,94 @@ _server = None _shutdown_event = None +_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist" + + +def _iter_frontend_fallback_candidates() -> "list[Path]": + """Yield `studio/frontend/dist` paths to try when the default is missing. + + Covers PATH-shadowed binaries whose __file__ resolves into a + site-packages tree that never received a vite build (e.g. plain + `pip install unsloth` from PyPI). + """ + import ast + import re + + out: list[Path] = [] + home_str = ( + os.environ.get("UNSLOTH_STUDIO_HOME") + or os.environ.get("STUDIO_HOME") + or str(Path.home() / ".unsloth" / "studio") + ) + venv_dir = Path(home_str).expanduser() / "unsloth_studio" + # Installer venv site-packages. + for pattern in ( + "lib/python*/site-packages/studio/frontend/dist", + "Lib/site-packages/studio/frontend/dist", + ): + out.extend(venv_dir.glob(pattern)) + # Editable source roots referenced from the installer venv. + for sp_pattern in ("lib/python*/site-packages", "Lib/site-packages"): + for sp in venv_dir.glob(sp_pattern): + for finder in sp.glob("__editable___*_finder.py"): + try: + src = finder.read_text(encoding = "utf-8") + except OSError: + continue + # Tolerate single- or multi-line dict literals; [^}]* still + # rejects nested dicts, which the setuptools template never + # emits for editable installs. + m = re.search( + r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S + ) + if not m: + continue + try: + mapping = ast.literal_eval(m.group(1)) + except (SyntaxError, ValueError): + continue + # Defensive: literal_eval can return a set / list / None if the + # matched literal is not a dict (regex captures `{...}`). + if not isinstance(mapping, dict): + continue + studio_pkg = mapping.get("studio") + if studio_pkg: + out.append(Path(studio_pkg) / "frontend" / "dist") + return out + + +def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Path]]: + """Pick a frontend dir that actually contains `index.html`. + + Returns (chosen, attempted). `chosen` is None if nothing servable was + found; `attempted` is the full ordered list for diagnostics. + """ + attempted: list[Path] = [] + seen: set[Path] = set() + + def _try(p: Path) -> bool: + try: + key = p.resolve() + except OSError: + key = p + if key in seen: + return False + seen.add(key) + attempted.append(p) + return (p / "index.html").is_file() + + if _try(Path(frontend_path)): + return attempted[-1], attempted + for alt in _iter_frontend_fallback_candidates(): + if _try(alt): + return attempted[-1], attempted + return None, attempted + + def run_server( host: str = "127.0.0.1", port: int = 8888, - frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist", + frontend_path: Path = _DEFAULT_FRONTEND_PATH, silent: bool = False, api_only: bool = False, llama_parallel_slots: int = 1, @@ -584,14 +669,48 @@ def run_server( print("=" * 50) print("") - # Setup frontend if path provided (skip in api-only mode) + # Setup frontend if path provided (skip in api-only mode). + # Falls back through alternate locations if the default lacks a built + # dist; errors out loudly rather than silently serving 404 on `/`. if frontend_path and not api_only: - if setup_frontend(app, frontend_path): + chosen, attempted = _resolve_frontend_path(Path(frontend_path)) + if chosen is not None and setup_frontend(app, chosen): if not silent: - print(f"[OK] Frontend loaded from {frontend_path}") + # Resolve so logs always show an absolute path for support. + try: + display = chosen.resolve() + except OSError: + display = chosen + print(f"[OK] Frontend loaded from {display}") else: - if not silent: - print(f"[WARNING] Frontend not found at {frontend_path}") + home_str = ( + os.environ.get("UNSLOTH_STUDIO_HOME") + or os.environ.get("STUDIO_HOME") + or str(Path.home() / ".unsloth" / "studio") + ) + # Windows ships the user-facing shim at $STUDIO_HOME/bin/unsloth.exe + # (a hardlink to the venv exe); Linux/macOS use the venv binary + # at $STUDIO_HOME/unsloth_studio/bin/unsloth. + home = Path(home_str).expanduser() + if sys.platform == "win32": + installer_bin = home / "bin" / "unsloth.exe" + else: + installer_bin = home / "unsloth_studio" / "bin" / "unsloth" + tried_lines = "\n".join(f" - {p}" for p in attempted) or " (none)" + raise SystemExit( + "[ERROR] Studio frontend build not found.\n" + f"Tried:\n{tried_lines}\n" + "\n" + "Likely cause: another 'unsloth' on PATH is shadowing the " + "installer's binary and points at a site-packages tree with " + "no built dist.\n" + "\n" + "Fix one of:\n" + f" - run the installer's binary directly: {installer_bin} studio\n" + " - pass --frontend \n" + " - pass --api-only to skip serving the web UI\n" + " - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh" + ) # Resolve once; shared by the log rewrite and the banner. display_host = _resolve_external_ip() if host == "0.0.0.0" else host @@ -718,7 +837,7 @@ if __name__ == "__main__": parser.add_argument( "--frontend", type = str, - default = Path(__file__).resolve().parent.parent / "frontend" / "dist", + default = _DEFAULT_FRONTEND_PATH, help = "Path to frontend build", ) parser.add_argument("--silent", action = "store_true", help = "Suppress output") diff --git a/studio/backend/tests/test_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py new file mode 100644 index 0000000000..2b49763386 --- /dev/null +++ b/studio/backend/tests/test_frontend_resolution.py @@ -0,0 +1,248 @@ +# 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 frontend-dist resolver in studio/backend/run.py. + +Loads only the relevant helpers via importlib so the test does not pull in +uvicorn / FastAPI / unsloth's full dependency tree. Pairs with the AST-style +test_host_defaults.py. +""" + +import ast +import importlib.util +import os +import sys +from pathlib import Path + +_RUN_PY = Path(__file__).resolve().parent.parent / "run.py" +_REPO_STUDIO_DIR = _RUN_PY.parent.parent # studio/ + + +def _load_helpers_only(): + """Import just the resolver helpers from run.py without executing the + server-side imports (uvicorn, structlog, etc.).""" + source = _RUN_PY.read_text(encoding = "utf-8") + tree = ast.parse(source) + keep = [] + wanted = { + "_DEFAULT_FRONTEND_PATH", + "_iter_frontend_fallback_candidates", + "_resolve_frontend_path", + } + for node in tree.body: + if isinstance(node, (ast.Import, ast.ImportFrom)): + keep.append(node) + elif isinstance(node, ast.Assign): + names = {t.id for t in node.targets if isinstance(t, ast.Name)} + if names & wanted: + keep.append(node) + elif isinstance(node, ast.FunctionDef) and node.name in wanted: + keep.append(node) + module = ast.Module(body = keep, type_ignores = []) + code = compile(module, str(_RUN_PY), "exec") + ns: dict = {"__file__": str(_RUN_PY), "__name__": "_run_helpers_test"} + exec(code, ns) + return ns + + +def test_resolver_returns_none_when_nothing_exists(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "no_studio")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](tmp_path / "missing") + assert chosen is None + assert attempted == [tmp_path / "missing"] + + +def test_resolver_picks_first_existing_candidate(tmp_path, monkeypatch): + dist = tmp_path / "good" / "frontend" / "dist" + dist.mkdir(parents = True) + (dist / "index.html").write_text("", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "no_studio")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](dist) + assert chosen == dist + assert attempted[-1] == dist + + +def test_resolver_falls_back_to_studio_home_site_packages(tmp_path, monkeypatch): + studio_home = tmp_path / "studio_home" + sp_dist = ( + studio_home + / "unsloth_studio" + / "lib" + / "python3.13" + / "site-packages" + / "studio" + / "frontend" + / "dist" + ) + sp_dist.mkdir(parents = True) + (sp_dist / "index.html").write_text("", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == sp_dist.resolve() + assert (tmp_path / "bogus") in attempted + + +def test_resolver_falls_back_via_editable_pth(tmp_path, monkeypatch): + """Simulates a `--local` install: dedicated venv with an editable .pth + pointing at a cloned repo that owns the built dist.""" + studio_home = tmp_path / "studio_home" + sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" + sp.mkdir(parents = True) + repo_root = tmp_path / "clone" + repo_studio = repo_root / "studio" + repo_dist = repo_studio / "frontend" / "dist" + repo_dist.mkdir(parents = True) + (repo_dist / "index.html").write_text("", encoding = "utf-8") + # Minimal `__editable___pkg_finder.py` carrying a MAPPING dict that + # setuptools' editable install generator writes. + finder = sp / "__editable___unsloth_0_0_0_finder.py" + finder.write_text( + "MAPPING: dict[str, str] = " + f"{{'studio': {str(repo_studio)!r}, 'unsloth': '/x', 'unsloth_cli': '/y'}}\n", + encoding = "utf-8", + ) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == repo_dist.resolve() + + +def test_iter_candidates_handles_missing_studio_home(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "nonexistent")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + # Glob over a non-existent dir is empty; must not raise. + candidates = helpers["_iter_frontend_fallback_candidates"]() + assert candidates == [] + + +def test_resolver_falls_back_to_windows_layout_site_packages(tmp_path, monkeypatch): + """Pins the `Lib/site-packages` (capital L) Windows venv layout + alongside the POSIX `lib/python*/site-packages` path.""" + studio_home = tmp_path / "studio_home" + sp_dist = ( + studio_home + / "unsloth_studio" + / "Lib" + / "site-packages" + / "studio" + / "frontend" + / "dist" + ) + sp_dist.mkdir(parents = True) + (sp_dist / "index.html").write_text("", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, _ = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == sp_dist.resolve() + + +def test_resolver_does_not_crash_on_non_dict_mapping_literal(tmp_path, monkeypatch): + """A finder file whose MAPPING value is a set / list / non-dict literal + (theoretically possible if the regex matched a brace-delimited literal + that ast.literal_eval can parse) must not AttributeError. The resolver + should skip that finder and keep probing.""" + studio_home = tmp_path / "studio_home" + sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" + sp.mkdir(parents = True) + # Bad finder: set literal, not a dict. ast.literal_eval parses it as set; + # any .get() call on it would raise AttributeError. + (sp / "__editable___bad_0_0_0_finder.py").write_text( + "MAPPING: dict[str, str] = {'studio', 'unsloth', 'unsloth_cli'}\n", + encoding = "utf-8", + ) + # Good finder that should still be discovered after the bad one is skipped. + repo_root = tmp_path / "clone" + repo_dist = repo_root / "studio" / "frontend" / "dist" + repo_dist.mkdir(parents = True) + (repo_dist / "index.html").write_text("", encoding = "utf-8") + (sp / "__editable___good_0_0_0_finder.py").write_text( + f"MAPPING: dict[str, str] = {{'studio': {str(repo_root / 'studio')!r}}}\n", + encoding = "utf-8", + ) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, _ = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == repo_dist.resolve() + + +def test_resolver_handles_multiline_mapping_dict(tmp_path, monkeypatch): + """A future setuptools / black reformat that wraps the MAPPING dict + across multiple lines must still parse and resolve. Locks in the + `[^}]*` + re.DOTALL behaviour.""" + studio_home = tmp_path / "studio_home" + sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" + sp.mkdir(parents = True) + repo_root = tmp_path / "clone" + repo_studio = repo_root / "studio" + repo_dist = repo_studio / "frontend" / "dist" + repo_dist.mkdir(parents = True) + (repo_dist / "index.html").write_text("", encoding = "utf-8") + finder = sp / "__editable___unsloth_0_0_0_finder.py" + finder.write_text( + "MAPPING: dict[str, str] = {\n" + f" 'studio': {str(repo_studio)!r},\n" + " 'unsloth': '/x',\n" + " 'unsloth_cli': '/y',\n" + "}\n", + encoding = "utf-8", + ) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, _ = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == repo_dist.resolve() + + +def test_systemexit_message_contains_actionable_fixes(tmp_path, monkeypatch): + """The user-facing recovery message is a contract: it must surface the + attempted paths and every concrete fix. Pin its structure so a future + refactor doesn't drop one.""" + import os + import sys + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "no_studio")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + bogus = tmp_path / "no_such_dist" + _, attempted = helpers["_resolve_frontend_path"](bogus) + home = Path(os.environ["UNSLOTH_STUDIO_HOME"]).expanduser() + if sys.platform == "win32": + installer_bin = home / "bin" / "unsloth.exe" + else: + installer_bin = home / "unsloth_studio" / "bin" / "unsloth" + tried_lines = "\n".join(f" - {p}" for p in attempted) + message = ( + "[ERROR] Studio frontend build not found.\n" + f"Tried:\n{tried_lines}\n" + "\n" + "Likely cause: another 'unsloth' on PATH is shadowing the " + "installer's binary and points at a site-packages tree with " + "no built dist.\n" + "\n" + "Fix one of:\n" + f" - run the installer's binary directly: {installer_bin} studio\n" + " - pass --frontend \n" + " - pass --api-only to skip serving the web UI\n" + " - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh" + ) + assert str(bogus) in message + assert "--frontend" in message + assert "--api-only" in message + assert "reinstall" in message + assert "installer's binary directly" in message + assert str(installer_bin) in message diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 67395a8378..e37cd0a8d8 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -206,6 +206,79 @@ def _find_setup_script() -> Optional[Path]: return None +def _iter_editable_studio_source_roots(venv_dir: Path): + """Yield repo roots from setuptools `__editable___*_finder.py` files in + *venv_dir*'s site-packages whose MAPPING includes a `studio` entry. + + Returns the parent dir of the mapped `studio` package (i.e. the repo + root), so callers can append `/studio/...` to reach any subdir. + """ + import ast + import re + + for sp_pattern in ("lib/python*/site-packages", "Lib/site-packages"): + for sp in venv_dir.glob(sp_pattern): + for finder in sp.glob("__editable___*_finder.py"): + try: + src = finder.read_text(encoding = "utf-8") + except OSError: + continue + # Tolerate single- or multi-line dict literals; [^}]* still + # rejects nested dicts, which the setuptools template never + # emits for editable installs. + m = re.search( + r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S + ) + if not m: + continue + try: + mapping = ast.literal_eval(m.group(1)) + except (SyntaxError, ValueError): + continue + # Defensive: literal_eval can return a set / list / None if the + # matched literal is not a dict (regex captures `{...}`). + if not isinstance(mapping, dict): + continue + studio_pkg = mapping.get("studio") + if studio_pkg: + yield Path(studio_pkg).parent + + +def _find_frontend_dist() -> Optional[Path]: + """Locate a built `studio/frontend/dist` (containing index.html). + + Probes (in order): package-local default, installer venv site-packages, + editable source roots referenced from the installer venv. Returns None + if nothing servable is found, so callers can decide to error or proceed + in `--api-only` mode. + + Fixes the silent 404 when another `unsloth` on PATH shadows the + installer's binary and points `_PACKAGE_ROOT` at a site-packages copy + that never received a vite build. + """ + candidates: List[Path] = [_PACKAGE_ROOT / "studio" / "frontend" / "dist"] + venv_dir = STUDIO_HOME / "unsloth_studio" + for pattern in ( + "lib/python*/site-packages/studio/frontend/dist", + "Lib/site-packages/studio/frontend/dist", + ): + candidates.extend(venv_dir.glob(pattern)) + for repo_root in _iter_editable_studio_source_roots(venv_dir): + candidates.append(repo_root / "studio" / "frontend" / "dist") + seen: set[Path] = set() + for c in candidates: + try: + resolved = c.resolve() + except OSError: + resolved = c + if resolved in seen: + continue + seen.add(resolved) + if (c / "index.html").is_file(): + return c + return None + + # ── helpers for `unsloth studio run` ──────────────────────────────── @@ -539,8 +612,14 @@ def studio_default( "--port", str(port), ] - if frontend: - args.extend(["--frontend", str(frontend)]) + # Resolve frontend explicitly so the spawned run.py uses a real + # built dist regardless of where its __file__ lands. Skip in + # --api-only (no UI served). + resolved_frontend = frontend + if resolved_frontend is None and not api_only: + resolved_frontend = _find_frontend_dist() + if resolved_frontend is not None: + args.extend(["--frontend", str(resolved_frontend)]) if silent: args.append("--silent") if api_only: From eacff8b8279c801fb7cf033f03de7f7b1e6bed63 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 26 May 2026 09:30:22 -0300 Subject: [PATCH 22/43] studio/frontend: show "1 second" instead of "1 seconds" in thinking blocks (#5777) * studio/frontend: show "1 second" instead of "1 seconds" in thinking blocks * Update studio/frontend/src/components/assistant-ui/reasoning.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- studio/frontend/src/components/assistant-ui/reasoning.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index e4401cc12e..285b87d138 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -134,7 +134,7 @@ function ReasoningTrigger({ {active ? ( Thinking... ) : ( - Thought for {duration ?? 0} seconds + Thought for {duration ?? 0} {duration === 1 ? "second" : "seconds"} )} Date: Tue, 26 May 2026 09:30:52 -0300 Subject: [PATCH 23/43] Studio: stop the model from replying twice when it refuses (#5775) --- studio/backend/core/inference/llama_cpp.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bf8a3c04df..76234386aa 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -60,7 +60,9 @@ _INTENT_SIGNAL = re.compile( # Handles both straight and curly apostrophes. # Excludes "I can", "I should", "I want to", "let's" which # appear frequently in direct answers / explanations. - r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b" + # Negative lookahead drops negated forms ("I will not", "I'll never") + # so a refusal doesn't trigger a re-prompt. + r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" r"|" # Step/plan framing: "First ...", "Step 1:", "Here's my plan" r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" From 953c8bfa4595b20fc80620048712d08a36a846bb Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 26 May 2026 10:02:03 -0300 Subject: [PATCH 24/43] studio/frontend: add space above the scroll-to-bottom arrow in chat (#5776) --- studio/frontend/src/components/assistant-ui/thread.tsx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index da243f37e1..a5bb4029d1 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -175,7 +175,8 @@ export const Thread: FC<{ From 849da89605aac35f94fd131fa2786be7ed7d78a0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 07:23:13 -0700 Subject: [PATCH 25/43] Fix unsloth studio update silently downgrading on macOS arm64 (#5767) * Fix unsloth studio update silently downgrading on macOS arm64 Root cause: studio/install_python_stack.py's "Updating base packages" step passes `--upgrade-package unsloth -r base.txt -c constraints.txt` with base.txt's `unsloth` and `unsloth-zoo` entries unpinned. On macOS arm64 the resolver silently backtracks to an older unsloth (2026.5.2 or even 2025.7.2) whenever a transitive constraint (the most common one is bitsandbytes wheel availability: 0.49.0+ ships macosx_14_0_arm64 wheels, older versions do not) makes the unpinned requirement satisfiable by an older release. install.sh already maintains an explicit `unsloth>=N.N.N` floor for the same reason, but the floor was missing from the in-venv update path. Reproduced on macos-14 across 2026.3.18 / 2026.4.8 / 2026.5.2 / 2026.5.6 starting states. All four ended on unsloth==2026.5.2 after a clean `unsloth studio update` invocation (2026.5.6 was a true downgrade, others were stale or partial advances). Fix mirrors install.sh: query PyPI at runtime for the current latest version of unsloth and unsloth-zoo, then pass `unsloth>=` and `unsloth-zoo>=` as extra positional pins alongside the existing `--upgrade-package` flags. Network failures fall back to the historical unpinned behaviour so offline installs continue to work. Applied to all three upgrade branches (standard update, local-repo overlay, no-torch). Also fix the cosmetic `Hardware detected: MLX -- Apple Silicon (i386)` banner. platform.processor() reads `uname -p` which returns "i386" on many universal2-shaped Python builds even on a native arm64 interpreter; platform.machine() is the reliable source ("arm64" once is_apple_silicon has gated us). * Dedup floor-pin call sites + LRU cache PyPI lookup Three upgrade branches each rebuilt the same conditional `unsloth>=` / `unsloth-zoo>=` arg list with two PyPI round-trips per branch -- six round-trips per `unsloth studio update` invocation. Extract a `_pin_floor_args(*, include_unsloth=True)` helper and wrap `_resolve_latest_pypi_version` in `functools.lru_cache` so the three branches share a single PyPI request per package. Functionally equivalent; pure cleanup on top of the previous commit. * Warn when PyPI is unreachable so the silent fallback is visible If `_resolve_latest_pypi_version` returns None for either lookup the floor args are silently dropped, which restores the pre-fix resolver behaviour. Print a single cyan `warning` line in `_pin_floor_args` when that happens so users behind a proxy / captive portal / firewalled PyPI mirror know the upgrade has degraded -- and can supply network egress or a `--index-url` mirror and retry. * Soft floor with unpinned-fallback for hosts where floor is unsatisfiable Reviewer found that the unconditional unsloth-zoo>=LATEST floor turns a previously-resolvable macOS 13 arm64 update into a hard resolver failure: unsloth-zoo 2026.5.4 requires mlx-vlm>=0.4.4 -> mlx>=0.30.0, and mlx 0.30+ only publishes macosx_14_0_arm64 wheels. The pre-fix behaviour backtracked to an older unsloth instead of erroring. We should not turn "stale" into "fail". Add pip_install_with_floor_fallback: first try the install with the floor appended; if the resolver cannot satisfy it (subprocess exit code != 0), retry the install without the floor and print a clear warning. The fall-through preserves the legacy "succeed-but-stale" contract on hosts where wheel availability is the bottleneck. Also extend pip_install_try with a req= kwarg so the floor attempt can pass `-r base.txt` like pip_install does, and add an UNSLOTH_NO_PYPI_FLOOR=1 opt-out for air-gapped CI / corporate PyPI mirrors that intentionally do not expose pypi.org directly. All three upgrade branches (standard, local-repo, no-torch) now go through the helper so the fallback behaviour is consistent. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add second fallback level: floor without constraints macOS arm64 floored attempt with -c constraints.txt fails because the single-env constraint `transformers==4.57.6` conflicts with the new unsloth-zoo 2026.5.4 -> mlx-vlm 0.4.4+ -> transformers>=5.1.0 chain. First fallback level retries the floored install without constraints (transformers freely resolves to a mlx-vlm-compatible version); downstream pip_install calls still apply constraints.txt to anything that doesn't transitively conflict. If THAT still fails (wheel availability rather than constraint conflict), drop the floor and fall back unpinned as before. Verified locally with uv pip compile against aarch64-apple-darwin python-3.13: strict-constrained floor errors, no-constraint floor resolves cleanly to unsloth==2026.5.7 + unsloth-zoo==2026.5.4 + transformers==5.5.0 + mlx-vlm==0.5.0. * setup.sh/.ps1: also gate fast-path on unsloth-zoo being up to date The version-check fast-path in setup.sh / setup.ps1 only looked at unsloth itself. If unsloth was at the PyPI latest but unsloth-zoo was stale, the gate set _SKIP_PYTHON_DEPS=true and install_python_stack.py never ran -- so the new floor pin from PR #5767 had no effect for the exact "unsloth at latest, zoo behind" state several reviewers flagged. Probe both packages' installed-vs-latest versions and only skip the deps step when BOTH match. When either is behind, fall through to install_python_stack.py so the new resolver fix gets a chance to run. Verified setup.sh with `bash -n`; the setup.ps1 change uses PowerShell if-expressions for the null-default pattern rather than bash-style ${var:-default} which is not valid PowerShell. * Skip unsloth-zoo floor too for custom no-torch test packages Reviewer found the asymmetric guard: the no-torch branch was already gating the unsloth floor on package_name == "unsloth" (test side packages may not publish to PyPI), but the unsloth-zoo floor was still added unconditionally. A custom no-torch update that ships its own forked zoo metadata could now hit a public PyPI floor that does not match the fork's published version. Add a symmetric `include_zoo` parameter to `_pin_floor_args` and gate both pins on the same `package_name == "unsloth"` check. * Address review feedback: simpler except clause + private-index note Gemini flagged TimeoutError in the PyPI fetch exception list. OSError already covers socket timeouts and the 3.11+ TimeoutError subclass on every supported Python, so drop the redundant entry and explain what each remaining exception catches. Codex flagged that floor lookups against pypi.org could break installs behind a lagging private mirror. Step 3 of pip_install_with_floor_fallback already recovers transparently in that case; expand the docstring so the behavior is discoverable without reading the body. * extras-no-deps: skip transformers==4.57.6 on macOS arm64 Reviewer flagged that the resolver-selected transformers from the no-constraints base step on macOS arm64 (transformers 5.x for mlx-vlm 0.4.4+) gets silently downgraded back to 4.57.6 by extras-no-deps.txt during the very next step, breaking mlx-vlm imports at runtime even though unsloth itself reports as latest. Add a PEP 508 platform marker so the pin only applies off macOS arm64. constraints.txt still enforces 4.57.6 everywhere else; mlx-vlm only publishes wheels for darwin arm64, so other platforms are unaffected. * setup.sh/.ps1: gate fast-path zoo probe on _PKG_NAME == unsloth Reviewer found the asymmetric custom-package regression: the new zoo-aware fast-path probes public unsloth-zoo unconditionally, but a custom STUDIO_PACKAGE_NAME side build may ship its own zoo fork via dependency metadata and not install public unsloth-zoo at all. The previous behaviour (skip Python deps if the custom package itself is at its declared latest) is preserved by only running the zoo probe when the managed package literally IS unsloth. Matches the include_zoo gate already in _pin_floor_args() at install_python_stack.py. * install_python_stack: all-or-nothing floor + uv-to-pip retry Two reviewer findings on the floor-pin helpers: 1. _pin_floor_args() previously kept a half-floor if one PyPI lookup succeeded and the other failed. With unsloth at latest but the zoo lookup down, the resolver could still backtrack zoo while we required unsloth at latest, defeating the pin. Return [] on any lookup failure so the unpinned legacy path runs cleanly. 2. pip_install_try() ran ONLY uv when USE_UV was true; a uv-specific failure short-circuited to False even when pip itself could have applied the floor. Mirror pip_install()'s uv-to-pip fallback: try uv, fall through to pip on non-zero exit, and only then give up. * extras-no-deps: rewrite marker without `not` for PEP 508 parsers pip's vendored packaging rejects `not (...)` in PEP 508 markers; the grammar only specifies `and` / `or` between boolean atoms. The staging macos-14 matrix failed every job at "Installing extras (no-deps)" with `Expected a marker variable or quoted string`. Apply De Morgan's law so the marker uses `or` between two `!=` checks, which both pip and uv parse cleanly. Behaviour identical: skip the 4.57.6 pin only on darwin arm64; pin everywhere else. * constraints: skip transformers==4.57.6 pin on macOS arm64 too Marker-gating the extras-no-deps.txt pin was not sufficient. Every subsequent pip_install in the update pipeline passes -c single-env/constraints.txt, and constraints.txt itself pinned transformers==4.57.6 unconditionally. The latest staging-2 run shows the base step's no-constraints fallback installed transformers 5.5.0 correctly, but a later constrained step (extras / studio / data-designer deps) silently downgraded it back to 4.57.6, leaving mlx-vlm 0.5.0 in the venv with an unsatisfied transformers>=5.5.0 requirement. Apply the same `sys_platform != "darwin" or platform_machine != "arm64"` marker to the constraints.txt entry so it is inert on darwin arm64. Other platforms still pin 4.57.6 because mlx-vlm only publishes wheels for darwin arm64; no other platform is affected. * constraints: carve out darwin arm64 from every == pin Marker-gating only transformers was not enough; staging-2 still failed with the same `transformers==4.57.6 in venv after the update` outcome because the resolver hit a `huggingface-hub==0.36.2` (and adjacent) conflict with mlx-vlm's `huggingface-hub>=1.5.0` requirement, then fell back to a stale stack even after my no-constraints level fired on the base step. Apply the same `sys_platform != "darwin" or platform_machine != "arm64"` marker to every == pin in constraints.txt. Range pins (mcp, fastmcp, websockets) stay active everywhere because they do not conflict with the mlx-vlm chain. mlx-vlm only publishes wheels for darwin arm64, so no other platform is affected. * install_python_stack: also --upgrade-package transformers and mlx-vlm Staging-2 showed that even after the constraints.txt carve-out for darwin arm64, the venv still ended up with the OLD `transformers==4.57.6` paired with a NEW `mlx-vlm==0.5.0` from unsloth-zoo's transitive upgrade. The resolver's --upgrade-package flag only freshens the named packages and their newly-pulled transitive deps; transformers was already installed at a version that satisfied unsloth-zoo's range (`>=4.51.3,<=5.5.0` with exclusions), so the resolver did not upgrade it -- even though mlx-vlm 0.5.0 requires `transformers>=5.5.0`. Add `--upgrade-package transformers` and `--upgrade-package mlx-vlm` to all three base-step branches. Both are no-ops when the package is absent (mlx-vlm only ships wheels on darwin arm64); on darwin arm64 this is what nudges the resolver to upgrade both together so the final venv is internally consistent. On Linux/Windows, transformers stays at 4.57.6 because constraints.txt still pins it there and mlx-vlm never enters the resolution. * install_python_stack: explicit mlx-vlm + transformers realign on macOS arm64 Even with --upgrade-package hints, uv leaves the venv with the already-installed transformers (4.57.6 inherited from the OLD venv's constrained install) when that version still happens to satisfy unsloth's own metadata range -- but it does not also re-resolve mlx-vlm's stricter `transformers>=5.5.0` requirement, so the venv ends up with mlx-vlm 0.5.0 paired with transformers 4.57.6 and mlx-vlm imports break at runtime. After the base step, on darwin arm64 only, run an explicit `pip install --upgrade mlx-vlm transformers` with constrain=False. This forces both packages through the resolver again as direct top-level requirements, so transformers is pulled up to whatever mlx-vlm's metadata requires (5.5.0 today). No effect on any other platform because mlx-vlm has no wheels off darwin arm64 and the branch is gated on IS_MAC_ARM. * requirements: marker-gate every == pin that conflicts with mlx-vlm chain Staging-2 kept ending up with transformers==4.57.6 even after the realign step, because studio.txt unconditionally pins huggingface-hub==0.36.2 (and datasets==4.3.0). Installing studio.txt with constraints active pulls the resolver back to a huggingface-hub that only recent transformers (4.x) supports, which silently downgrades the realigned 5.5.0 to 4.57.6 -- exactly the inconsistency we tried to prevent. Also extras-no-deps.txt still pinned trl==0.23.1 unconditionally; the 0.23.1 wheel transitively requires huggingface-hub<1, same coupling. Marker-gate all three. The carve-out is identical to constraints.txt's: inactive on darwin arm64 (where the mlx-vlm chain dictates newer versions), active everywhere else (where Linux/Windows users rely on the single-env pins). mlx-vlm only publishes wheels for darwin arm64 so no other platform is affected. * realign: --force-reinstall mlx-vlm + transformers + huggingface_hub Plain --upgrade does not force uv to re-resolve mlx-vlm's transformers requirement when the already-installed transformers happens to satisfy unsloth's own range. Switch to --force-reinstall on the three packages so the resolver tears them down and brings them back together with consistent versions. Include huggingface_hub because transformers 5.x requires hf-hub>=1.5.0 and the resolver would not touch it otherwise. * realign: pin transformers via mlx-vlm's own metadata spec `pip install --force-reinstall mlx-vlm transformers` still resolved to an already-installed transformers 4.57.6 because uv treats it as satisfying unsloth's transformers range without re-checking mlx-vlm's stricter requirement. Pull mlx-vlm's actual transformers specifier from its installed metadata at runtime and pass it as an explicit version requirement (e.g. `transformers>=5.5.0` for mlx-vlm 0.5.0). That removes the resolver's wiggle room: it MUST pick a transformers satisfying mlx-vlm AND unsloth, which on darwin arm64 with the latest unsloth-zoo means transformers==5.5.0. Falls back to unpinned `transformers` if metadata read fails, so this never errors. * realign: uninstall-then-install to bypass uv's incumbent bias Every flag-based approach failed: --upgrade, --upgrade-package, --force-reinstall, and even an explicit `transformers>=5.5.0` requirement all left the venv with transformers==4.57.6 because uv treats the already-installed version as satisfying unsloth-zoo's range and refuses to disturb it, even when it does not satisfy mlx-vlm's stricter requirement. Replace the realign step with an explicit uninstall of the conflicting trio (transformers / mlx-vlm / huggingface_hub) followed by a fresh install. With no transformers in the venv, the resolver MUST pick a version satisfying every installed package's metadata, which on darwin arm64 with the latest unsloth-zoo is uniquely 5.5.0. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim verbose comments across PR #5767 changes * Simplify mac-arm64 fix: install MLX stack with --no-deps The previous approach (PyPI floor pin + 3-level fallback + macOS arm64 realign step + marker carve-outs on every == pin) was fighting symptoms. The root cause is that unsloth-zoo declares mlx-vlm>=0.4.4 as a darwin arm64 dep, and mlx-vlm 0.5.0's metadata pulls in transformers>=5.5.0, which conflicts with the main venv's transformers==4.57.6 pin and forces the resolver to backtrack unsloth. Severing that chain at its source: install mlx + mlx-metal + mlx-lm + mlx-vlm with --no-deps BEFORE unsloth-zoo. The resolver sees mlx-vlm already installed (>=0.4.4) and never inspects its transformers metadata. Per-model transformers version routing is already handled at runtime by the side-car venvs in utils/transformers_version.py (.venv_t5_530 for Ministral/GLM/Qwen3 MoE, .venv_t5_550 for Gemma 4). Net change: -224 / +71 lines across install.sh, install_python_stack.py and the three requirements files. Reverted: - _resolve_latest_pypi_version + _pin_floor_args + pip_install_with_floor_fallback - macOS arm64 realign step (pip uninstall + reinstall) - --upgrade-package transformers --upgrade-package mlx-vlm in base steps - All ; sys_platform != "darwin" or platform_machine != "arm64" markers in constraints.txt, studio.txt, extras-no-deps.txt - pip_install_try restored to its pre-PR signature Added: - install.sh: Apple Silicon MLX --no-deps install before unsloth (both fresh and migrated branches) - install_python_stack.py: same step gated on IS_MAC_ARM and not skip_base Kept (independent bugs): - setup.sh / setup.ps1 dual-package zoo version check - platform.processor() -> platform.machine() hardware-detect fix * Minimise PR to mac-arm64-specific changes only Revert setup.sh and setup.ps1 to main -- the dual-package zoo check was defensive and not strictly needed once mlx-vlm is installed --no-deps (the resolver-backtrack scenario that produced stale zoo no longer happens). Tighten remaining comments in install.sh and install_python_stack.py. Final PR-attributable changes: install.sh +24/-5 (MLX --no-deps in 2 places) studio/install_python_stack.py +19 (MLX --no-deps + IS_MAC_ARM) studio/backend/utils/hardware/hardware.py +6/-6 (processor() -> machine()) studio/backend/requirements/*.txt unchanged * Revert "Minimise PR to mac-arm64-specific changes only" This reverts commit 9470daa855f8c1588350585c6d7b82041228e32e. * Revert "Simplify mac-arm64 fix: install MLX stack with --no-deps" This reverts commit f8a43b87e8b46fd5e3b7942d38f15b58f3844aae. * Revert "Trim verbose comments across PR #5767 changes" This reverts commit c3f293a10fa11a74212a6a7b97f46e619dedc38e. * Simplify mac-arm64 fix: --no-deps MLX + METADATA patch Root cause: unsloth-zoo declares mlx-vlm>=0.4.4 as a darwin-arm64 dep, and mlx-vlm 0.5.0's published metadata declares transformers>=5.5.0. Every subsequent resolver run with constraints.txt's transformers==4.57.6 sees the conflict and backtracks unsloth to escape it (user-reported downgrade). The aggressive pin doesn't reflect what mlx-vlm actually requires at top-level import time -- the symbols it loads (AutoProcessor, AutoTokenizer, ProcessorMixin, BatchFeature) are stable across transformers 4.51+. Model- specific submodules that genuinely need 5.x APIs are only loaded once the 3-tier transformers dispatcher (utils/transformers_version.py) has activated the matching .venv_t5_530 / .venv_t5_550 side-car at runtime. Fix: on Apple Silicon, install the MLX stack with --no-deps then rewrite mlx-vlm/mlx-lm's installed METADATA to declare transformers>=4.51.3. Now the resolver sees mlx-vlm 0.5.0 as compatible with the main venv's transformers==4.57.6 and there's nothing to backtrack. Reverts the previous heavy machinery: - _resolve_latest_pypi_version, _pin_floor_args, pip_install_with_floor_fallback - macOS arm64 realign step (pip uninstall + reinstall) - --upgrade-package transformers --upgrade-package mlx-vlm in base steps - All ; sys_platform != "darwin" or platform_machine != "arm64" markers in constraints.txt / studio.txt / extras-no-deps.txt - setup.sh / setup.ps1 dual-package zoo check (Windows never had the bug; with this fix in place stale zoo no longer happens on macOS either) - pip_install_try restored to pre-PR signature Kept: - install.sh: MLX --no-deps install in fresh + migrated branches - install_python_stack.py: same step gated on IS_MAC_ARM and not skip_base - _relax_mlx_metadata() helper, called immediately after each MLX install - studio/backend/utils/hardware/hardware.py: platform.processor() -> platform.machine() cosmetic fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use UV_OVERRIDE to relax mlx-vlm transformers pin uv supports --overrides / UV_OVERRIDE which globally overrides any package's stated dependency requirement. mlx-vlm 0.5.0 declares transformers>=5.5.0 and mlx-lm 0.31.3 declares transformers>=5.0.0; neither is true at top-level import time (their imports use AutoProcessor / AutoTokenizer / ProcessorMixin / BatchFeature which are stable across transformers 4.51+). Per-model 5.x routing is handled at runtime via the .venv_t5_530 / .venv_t5_550 side-cars. Override file (overrides-darwin-arm64.txt) declares transformers>=4.51.3 ; exported via UV_OVERRIDE env var on Apple Silicon by both install.sh and install_python_stack.py. uv then resolves mlx-vlm as compatible with the main venv's transformers==4.57.6 (constraints.txt) and unsloth advances cleanly to LATEST. Drops, vs. the previous attempts: - _resolve_latest_pypi_version + _pin_floor_args + pip_install_with_floor_fallback (floor-pin machinery -- replaced by single UV_OVERRIDE line) - macOS arm64 realign step (pip uninstall + reinstall) - --upgrade-package transformers --upgrade-package mlx-vlm in base steps - All ; sys_platform != "darwin" or platform_machine != "arm64" markers - _relax_mlx_metadata() helper + sed METADATA patch (uv reads from index, not dist-info, so dist-info patches were ineffective) Kept: - install.sh / install_python_stack.py: MLX latest install on Apple Silicon (now without --no-deps, the override lets the resolver pick a consistent set) - studio/backend/utils/hardware/hardware.py: platform.machine() cosmetic fix * Trim UV_OVERRIDE comments; bump override floor to 4.57.6 Match the main venv's constraints.txt pin exactly so the override file reads as the actual installed version rather than mlx-vlm's API floor. Comments collapsed to one-liners where possible. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.sh | 14 +++++++------ .../single-env/overrides-darwin-arm64.txt | 5 +++++ studio/backend/utils/hardware/hardware.py | 15 ++++++------- studio/install_python_stack.py | 21 +++++++++++++++++++ 4 files changed, 42 insertions(+), 13 deletions(-) create mode 100644 studio/backend/requirements/single-env/overrides-darwin-arm64.txt diff --git a/install.sh b/install.sh index d12abe298f..49bb7a7b89 100755 --- a/install.sh +++ b/install.sh @@ -1290,6 +1290,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi +# Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file). +if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" + if [ -f "$_OVERRIDES_FILE" ]; then + export UV_OVERRIDE="$_OVERRIDES_FILE" + fi +fi + _TAURI_INITIAL_GPU_BRANCH="unknown" if [ "$SKIP_TORCH" = true ]; then _TAURI_INITIAL_GPU_BRANCH="no_torch" @@ -2108,12 +2116,6 @@ else fi fi -# ── Install mlx-vlm on Apple Silicon (optional, for VLM training) ── -if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then - substep "installing mlx-vlm (VLM training support)..." - run_install_cmd "install mlx-vlm" uv pip install --python "$_VENV_PY" mlx-vlm -fi - # ── Run studio setup ── tauri_log "STEP" "Running Studio setup" # When --local, use the repo's own setup.sh directly. diff --git a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt new file mode 100644 index 0000000000..2cd03d8b78 --- /dev/null +++ b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt @@ -0,0 +1,5 @@ +# mlx-vlm / mlx-lm declare transformers>=5.x which conflicts with the +# main venv's constraints.txt pin transformers==4.57.6 and forces uv to +# backtrack unsloth. Relax to match the pin -- per-model 5.x routing +# happens at runtime via the side-car venvs. +transformers>=4.57.6 diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 3764e38272..ede37e2953 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -144,7 +144,10 @@ def detect_hardware() -> DeviceType: if is_apple_silicon() and _has_mlx(): DEVICE = DeviceType.MLX CHAT_ONLY = False - chip = platform.processor() or platform.machine() + # platform.processor() runs `uname -p` which returns "i386" on most + # universal2 / Rosetta-shaped Python builds even on native arm64. + # platform.machine() is "arm64" once is_apple_silicon() has gated us. + chip = platform.machine() or "arm64" print(f"Hardware detected: MLX — Apple Silicon ({chip})") return DEVICE @@ -279,13 +282,11 @@ def get_gpu_memory_info() -> Dict[str, Any]: try: info = mx.device_info() - gpu_name = ( - info.get("device_name") - or platform.processor() - or platform.machine() - ) + # See detect_hardware(): platform.processor() can return "i386" + # on native arm64 Python builds, so prefer machine() as fallback. + gpu_name = info.get("device_name") or platform.machine() or "arm64" except Exception: - gpu_name = platform.processor() or platform.machine() + gpu_name = platform.machine() or "arm64" return { "available": True, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9166d35ce3..ca7fe3f004 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -37,6 +37,7 @@ from backend.utils.wheel_utils import ( IS_WINDOWS = sys.platform == "win32" IS_MACOS = sys.platform == "darwin" IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64" +IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64" # ── ROCm / AMD GPU support ───────────────────────────────────────────────────── # Mapping from detected ROCm (major, minor) to the best PyTorch wheel tag on @@ -423,6 +424,7 @@ def _infer_no_torch() -> bool: NO_TORCH = _infer_no_torch() + # -- Verbosity control ---------------------------------------------------------- # By default the installer shows a minimal progress bar (one line, in-place). # Set UNSLOTH_VERBOSE=1 in the environment to restore full per-step output: @@ -448,6 +450,11 @@ LOCAL_DD_GITHUB_PLUGIN = ( SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" ) +# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides file). +_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt" +if IS_MAC_ARM and _MLX_OVERRIDES.is_file(): + os.environ.setdefault("UV_OVERRIDE", str(_MLX_OVERRIDES)) + # -- Unicode-safe printing --------------------------------------------- # On Windows the default console encoding can be a legacy code page # (e.g. CP1252) that cannot represent Unicode glyphs such as ✅ or ❌. @@ -960,6 +967,20 @@ def install_python_stack() -> int: [sys.executable, "-m", "pip", "install", "--upgrade", "pip"], ) + # macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the + # mlx-vlm / mlx-lm transformers pin -- set at module load). + if IS_MAC_ARM and not skip_base: + _progress("MLX stack (Apple Silicon)") + pip_install( + "Installing MLX stack (mlx + mlx-lm + mlx-vlm)", + "--no-cache-dir", + "--upgrade", + "mlx", + "mlx-metal", + "mlx-lm", + "mlx-vlm", + ) + # 3. Core packages: unsloth-zoo + unsloth (or custom package name) if skip_base: pass From e57a1a73b8949b20bd8f03ac1909d06dbe555116 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 07:25:14 -0700 Subject: [PATCH 26/43] Update _utils.py --- unsloth/models/_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b940fdf35a..8965a88fae 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.7" +__version__ = "2026.5.8" __all__ = [ "SUPPORTS_BFLOAT16", From 1cf145c070eaaf050b2f21f2a649ede0ed3e9862 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 07:36:10 -0700 Subject: [PATCH 27/43] Bump install.sh / install.ps1 pin to unsloth>=2026.5.8 (#5791) PyPI release unsloth 2026.5.8 is now live. Bumps the pinned floor in install.sh and install.ps1 from unsloth>=2026.5.7 to unsloth>=2026.5.8 so fresh installs resolve to the new wheel. --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 52766370d1..1ed8fafdb4 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1300,7 +1300,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.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1314,7 +1314,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.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1352,7 +1352,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.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1364,7 +1364,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1392,7 +1392,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.7" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.8" --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 49bb7a7b89..0548ee2a57 100755 --- a/install.sh +++ b/install.sh @@ -1873,7 +1873,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.7" unsloth-zoo + "unsloth>=2026.5.8" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -1886,7 +1886,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.7" unsloth-zoo + "unsloth>=2026.5.8" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2054,7 +2054,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.7" unsloth-zoo + "unsloth>=2026.5.8" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2072,7 +2072,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.7" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.8" 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..." @@ -2104,7 +2104,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.7" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.8" --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 649b9f780865cb110a2200a3199110f8049c343d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 23:13:45 -0700 Subject: [PATCH 28/43] Studio: expose --parallel / -np flag on `unsloth studio run` (#5737) * Studio: expose --parallel / -np on `unsloth studio run` The CLI was hardcoding `llama_parallel_slots=4` in `run_kwargs` at `unsloth_cli/commands/studio.py`, leaving users unable to tune the concurrent decode slot count even though the engine, KV-cache math, and `studio.backend.run.run_server(llama_parallel_slots=...)` plumbing all already accepted any N. This change adds a `--parallel` / `--n-parallel` / `-np` typer option (default 4 -- matches the previous hardcoded value), forwards it into `run_kwargs`, and pins the new surface with 4 unit tests. Per-request state in `routes/inference.py` is already isolated (`cancel_event` and `prev_text` are per-request locals in every streaming handler; the `_lock` / `_serial_load_lock` only wrap load/unload, not chat completions), so no concurrency refactor is needed alongside this -- the engine layer already handles N concurrent requests on one loaded model when llama-server is told to. Range guards: 1 <= N <= 64. With higher N each slot gets ctx/N KV cache; users tuning this should be aware that per-call context shrinks proportionally. `unsloth studio` (the bare default command, no subcommand) still defaults to llama_parallel_slots=1 via `run_server`'s own default; this PR does not change that path -- it only exposes the knob on the one-liner `studio run` command that already silently used 4. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Forward --parallel through venv re-exec and drop colliding short aliases `unsloth studio run` re-execs into the Studio venv when invoked from outside it (the common path). The arg-builder forwards every typer option but the new --parallel, so the child re-execs at the default 4 and any user value is silently dropped. Worse: pre-PR users who already pass `-np N` as a pass-through extra (where llama.cpp's last-wins parsing made it stick) silently lose N after this PR lands. Forward --parallel explicitly in the re-exec arg list. While auditing the re-exec path, also drop the colliding 1-char short aliases -m (--model) and -f (--frontend) plus the redundant -hfr. Click's short-option clustering had been silently mis-parsing ~11 llama-server short flags via the pass-through path: -fa as `-f a`, -mg 0 as `-m g` + stray 0, -fitt 1024 as `-f itt` + stray 1024, -hff path as `-f f` + stray `-h path`, -cmoe / -cram / -sm / -ncmoe etc. The docstring promise ("any flag this command does not recognize is forwarded verbatim") was silently violated. -hf (2-char) is kept because Click treats multi-char shorts atomically (no clustering of -hff / -hfv / -hffv / -hft) and -hf is documented in basics/api/README.md. --model / --hf-repo / --frontend long forms all unchanged. studio_default keeps -f because it has no pass-through. Tests: - test_studio_run_parallel_flag.py: 8 new re-exec coverage cases (all 3 aliases, 3 platforms via sys.platform mock, pre-PR `-np` regression, mixed with pass-through extras). - test_studio_run_short_alias_clashes.py (new): surface checks that the removed shorts cannot reappear, plus 11 parametrized cases proving each previously-broken llama-server short flag now passes through verbatim, plus a happy-path test that documented -hf still works for `org/repo:variant` syntax. All 27 tests pass. Negative test (revert either fix) shows the new tests catch the regression. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix stale studio run docstring describing rejected llama-server flags The pre-PR docstring listed --port, -c / --ctx-size, --api-key, -ngl, --jinja, --flash-attn, --no-context-shift as "rejected with HTTP 400", but only --port and --api-key (plus other networking / auth / model identity / single-model UI flags) are actually in studio/backend/core/inference/llama_server_args.py's denylist. -c / -ngl / --jinja / --flash-attn / --no-context-shift are pass-through and last-wins-override Studio's auto-set value. Rewrite the docstring to match the real denylist groups and point at the canonical source. Also add --parallel to one of the examples now that it is a first-class flag. * ci: broaden Linux + narrow Windows llama.cpp runtime patterns + trim #5741 comments (#5746) * ci: broaden Linux llama.cpp runtime pattern to lib*.so* #5741 patched the explicit Linux pattern list to add ``libllama-*-impl.so*`` after ggml-org/llama.cpp#23462 (between b9279 and b9283) split each binary's entry code into a paired ``lib-impl.so`` shared library. Same class of upstream repackaging will hit us again whenever a new shared lib is added. Mirror what macOS already does and replace the per-lib list with a single ``lib*.so*`` glob. ``copy_globs`` (line 3614) unions patterns, so the per-variant ``libggml-cuda.so*`` / ``libggml-hip.so*`` entries were never filtering anything; the spec lives in ``runtime_payload_health_groups`` (line 5209) which keeps the explicit minimum-required list per variant. Dry-run against b9296-bin-ubuntu-x64.tar.gz: 40 files copied (all ggml, llama, mtmd, impl variants + the two binaries we ship), 22 skipped (other CLIs, rpc-server, LICENSE). Functionally equal to the post-#5741 set. * cleanup: trim #5741 comments on the pydantic split Comments added in #5741 explained the original bug in full each time. They are mostly redundant with the commit message and the PR. Trim them to one short paragraph per site. No behavior change. * ci: narrow Windows runtime pattern to llama-server.exe + llama-quantize.exe Studio only invokes llama-server and llama-quantize. Mac and Linux already filter to those two binaries; Windows was the odd one out with ``*.exe`` copying every CLI upstream ships (llama-cli, llama-bench, llama-mtmd-cli, ...). Dry-run on b9296 (win cpu-x64, cpu-arm64, cuda-13.1, hip-radeon): 20 unused EXEs skipped per variant, all DLLs (incl. the new llama-*-impl.dll family) still copied via ``*.dll``. ``existing_install_matches_choice`` already checks llama-server.exe exists explicitly (line 5297), so the health gate is unchanged. * Lower default weight_decay in RL config from 0.01 to 0.001 (#5747) In full FT, AdamW weight decay shrinks the parameter directly so the implicit prior is W -> 0. In LoRA the trained parameters are A and B while the effective weight is W = W_init + (alpha/r) * B @ A; decaying A and B separately drives BA -> 0, hence W -> W_init rather than 0. The previous default of 0.01 inherited from full-FT recipes adds a measurable pull on the merged adapter back toward the base model over a few thousand steps. 0.001 keeps a small Frobenius-norm prior on ||A||^2 + ||B||^2 for numerical stability without meaningfully biasing the merged weight toward init, and aligns with the value used across the unsloth notebook templates. * Studio: strip orphan tool_call XML leaking into visible content (#5735) * Studio: strip orphan tool_call XML from streamed visible content The speculative-buffer state machine in `studio/backend/core/inference/llama_cpp.py` can slice a tool_call XML block between the silent DRAINING path and the user-visible content_accum, depending on when in the model's emission the BUFFERING -> STREAMING -> DRAINING transitions fire. Three leak shapes were observed in a 2026-05-22 sweep of 900 Qwen3.5 / Qwen3.6 GGUF runs: Pre-fix XML leak rate: 20/900 (2.22%), concentrated 6.7% on the larger Q8 / MTP configs: Qwen3.6-35B-A3B Q8_0 4/60 (6.7%) Qwen3.6-35B-A3B-MTP Q4 4/60 (6.7%) Qwen3.5-35B-A3B Q8_0 3/60 (5.0%) Qwen3.6-27B Q8_0 3/60 (5.0%) The existing `_TOOL_XML_RE` only matched well-formed `...` and `` pairs, so unterminated openings (close was DRAINED) and orphan closes (opening was DRAINED) survived the strip and reached the user. Fix relaxes the regex to also strip: 1. Orphan opening up to end-of-string: `(?:|\Z)` 2. Orphan closing tag: bare `` / `` Verified on the full sweep: 20/900 -> 0/900 (100% of detected leaks eliminated). 16 unit tests in `test_tool_xml_strip.py` pin all three leak shapes plus the well-formed cases, plus parametrised checks on the 5 actual real-world leak samples from the sweep data. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: strip tail-only orphan + tighten regex The 2026-05-22 gdpval sweep surfaced a 4th XML-leak shape not caught by the earlier regex: a bare `\n\n` at end-of-buffer (7 of 192 trials, all Qwen3.5-27B + a few Qwen3.6-27B). The model emits the full `...content... ` envelope, the speculative buffer DRAINS the opening tags as intended, but EOS (max_tokens cutoff) truncates the outer `` close, leaving just `` as the visible tail. We strip this ONLY when end-anchored (`\s*\Z`) so legitimate mid-text uses (user code samples, documentation discussing the Qwen tool-call XML shape) survive. Verified on the 192-trial gdpval corpus: before=7, after=0. While at it, fold the five top-level alternations into three by sharing tag-name and prefix subgroups: ... + ... + --> <(?:tool_call|function=\w+)>... | --> Semantically identical (verified by replay over the 192-trial corpus + adversarial inputs, 0 diffs) and 1.34x faster on real workloads. Backtracking-safety pinned by two new perf guards (256KB '<' spam, 1000x orphan opens). Tests: 16 -> 28 (6 new functional + 4 well-formed-vs-orphan + 2 perf guards). * Tighten comments in XML-strip regex and tests Code says what it does; comments were repeating it. Strip the verbose explanations down to the WHY-only bits (engine quirk, tail-anchor rationale, real-world source of each test sample). No code changes. inference.py: 21 -> 12 lines around _TOOL_XML_RE test_tool_xml_strip.py: 343 -> 259 lines (-84) Tests: 28/28 still pass. * [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> * Address review: deny pass-through --parallel, preserve legacy short aliases, fix test harness Round 1 review fixes for #5737: 1. Deny --parallel / --n-parallel / -np in the pass-through validator. Without this, `unsloth studio run --model X --parallel 8 -- --parallel 999` would last-win-override the running llama-server slot count while Studio's app.state.llama_parallel_slots and KV-cache fitting stay at the typer value (8), so the resource plan and the running process disagree. Also bypasses the typer 1..64 range guard. Reject so the only path is the first-class typer flag. 2. Backwards-compat shim for -m / -hfr / -f. Dropping the short aliases from typer broke any script using `unsloth studio run -m X` or `-hfr Y` or `-f dist`. Add _consume_legacy_short_aliases which pops EXACT whole-token matches (or `-x=value` inline form) from ctx.args into the corresponding typer parameter. Clustered tokens (`-fa`, `-mg`, `-fitt`, ...) are left in the pass-through tail unchanged. --model becomes Optional with an explicit missing-required check after the preprocessor so legacy `-m X` still satisfies the "must specify a model" requirement. 3. Drop mix_stderr from CliRunner. Typer 0.25.1 / Click 8.4.1 removed the kwarg; the test harness raised TypeError before exercising the PR behaviour. Tests run cleanly on current and older Typer/Click. 4. Correct the -np regression test docstring. Pre-PR `-np 8` was clustered by Click as `-p 8` (port=8) + stray `-n`, silently breaking the port binding -- not "passed through as 8 slots". The post-PR assertion (child gets --parallel 8) is unchanged. 5. Update studio run docstring listing rejected flags so it now correctly includes --parallel / -np / --n-parallel. New tests: - test_llama_server_args.py: parametrized denylist coverage for --parallel / --n-parallel / -np including equals-form, including out-of-range bypass attempts (999, 0). is_managed_flag flips True. - test_studio_run_short_alias_clashes.py: legacy -m / -hfr / -f promote to typer params; --model X + -m Y conflict errors; clustered -mg / -fa / -fitt still pass through (the original bug fix holds). 132 tests pass (98 backend + 34 cli). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Extend legacy-alias shim tests for repo:variant, inline value form, and missing model Three additional edge cases for the -m / -hfr / -f preprocessor: - `-m unsloth/foo:UD-Q4_K_XL` round-trips through both the preprocessor and _split_repo_variant so the child sees --model + --gguf-variant. - `-m=foo` inline value form is promoted just like `-m foo`. - Missing --model after the preprocessor raises typer.Exit(2) cleanly (replacing typer's pre-PR required-flag enforcement now that --model is Optional to allow the legacy promotion path). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scrub .github/workflows for staging push (matches staging base) * Fix studio CLI argv handling and pass-through docstring drift - studio/backend/core/inference/llama_server_args.py: drop the stale ``-np``/``--parallel`` entry from the docstring's pass-through tunable list. These flags moved into _DENYLIST_GROUPS so the docstring now contradicts the validator and would mislead future maintainers debugging the ValueError from validate_extra_args(["--parallel","8"]). The deleted wording was introduced by dbea77e34 ("Studio: forward llama-server args from `unsloth studio run`, activate `unsloth run`, and allow passing model:quant to load models") when --parallel was still a documented pass-through; the same commit's "quant" reference is about the model:quant syntax, unrelated to the parallel slot wording being deleted here. - unsloth_cli/commands/studio.py: add _expand_attached_np_short next to _consume_legacy_short_aliases. Both work around Click's short-option clustering for this command -- the legacy preprocessor for `-m` / `-f` / `-hfr` and this one for the attached `-np` form. Click clusters `-np8` as `-n -p 8` because `-p` is the typer short for `--port`, silently setting port=8 and dropping the parallel value; rewriting the attached form into separated `-np ` in sys.argv before Click parses preserves the user's value. Space/equals forms (`-np 8`, `-np=8`) already work and are left alone. - unsloth_cli/__init__.py: import _expand_attached_np_short from the studio command and run it only when argv[0] looks like the unsloth console-script or workspace cli.py, so importing this module from a notebook or pytest run does not mutate the caller's argv. * Tighten the -np canonicaliser comments Drop the helper's co-location sentence (location is self-evident from grep) and shorten the entry-gate rationale to one short sentence covering the why. * Sync .github/workflows with upstream author branch * Sync .github/workflows with upstream author branch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Bump install.sh / install.ps1 pin to unsloth>=2026.5.7 (#5753) PyPI release unsloth 2026.5.7 is now live. Bumps the pinned floor in install.sh and install.ps1 from unsloth>=2026.5.6 to unsloth>=2026.5.7 so fresh installs resolve to the new wheel. Tagged on main as v0.1.416-beta. * Catch attached `-np` form in backend pass-through validator The CLI-side `_expand_attached_np_short` rewrites `-np8` to `-np 8` before Click parses, but HTTP /load `llama_extra_args=["-np8"]` goes straight to `validate_extra_args` which only matched the exact token. Reproducer: `validate_extra_args(["-np8"])` previously returned `["-np8"]` instead of raising; once forwarded to llama-server it last-win-overrode Studio's slot count while `app.state.llama_parallel_slots` stayed at the typer value. Normalise `-np` to `-np` in `_flag_name` so the denylist catches the attached form alongside `-np`, `-np=8`, `--parallel`, `--parallel=8`, and `--n-parallel`. Tests parametrize the new form including out-of-range values. * Restore _consume_legacy_short_aliases unit tests + _expand_attached_np_short tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore .github/workflows from origin/main Earlier merge from claude_review's staging-scrub commits accidentally deleted production CI workflows. Restore them to main's state. * Scrub .github/workflows for staging push (matches staging base) * Sync .github/workflows with upstream author branch * Round 5+6: broaden -np gate to exact basenames + runtime parallel test Reviewer-flagged improvements squashed into one commit so the auto-push review bot doesn't keep stomping the branch: - unsloth_cli/__init__.py: exact-basename match instead of endswith('cli.py'). Covers unsloth, unsloth.exe, unsloth-cli, unsloth-cli.exe, cli.py, unsloth-cli.py. A third-party mycli.py that happens to import unsloth_cli no longer has its argv mutated. - unsloth_cli/tests/test_studio_run_parallel_flag.py: parametrised runtime test (N in {1, 4, 8, 64}) that fakes the in-venv path and asserts run_server is invoked with llama_parallel_slots=N. Complements the existing source-text check so refactors that preserve runtime semantics don't trip a false failure. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Round 7: respect '--' end-of-options and reject flag-as-value Round 7 reviewer flagged three legitimate edge cases: - _expand_attached_np_short rewrote post-'--' tokens. Convention: '--' ends option processing; payload after it is raw. Stop the loop there. - _consume_legacy_short_aliases promoted post-'--' legacy aliases for the same reason. Treat post-'--' tail as raw. - Legacy '-m -fa' silently consumed '-fa' as the model name, hiding the real CLI shape error. Reject any next-token that starts with '-' (except the lone '-' stdin/path sentinel) with a clear BadParameter. Also expanded the missing-model error string to mention the still- supported legacy '-m' / '-hfr' aliases so users hitting that diagnostic on legacy scripts get the right migration hint. Added four regression tests covering each new behaviour. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Round 8: soften flag-as-value to long-form only + normalise is_managed_flag Round 8 reviewer flagged two cleanups: - _consume_legacy_short_aliases rejected any next token starting with '-' as a flag, which would break legitimate values like '-foo' (path or model name with leading dash). Narrow the rejection to '--long' tokens only; '-x' short forms still pass through. - is_managed_flag did raw _DENYLIST membership while validate_extra_args goes through _flag_name first, so '-np8' / '--parallel=8' / '--port=9000' classified as not-managed by the helper but rejected by the validator. Route is_managed_flag through _flag_name so the two helpers agree on every form callers might use. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Round 9: also catch -np-1 / -np+1 signed attached forms in denylist Round 9 reviewer noticed _flag_name normalised -np but missed signed variants -np-1 and -np+1, so validate_extra_args waved them through while rejecting --parallel -1. llama.cpp would error out on negative slot counts anyway, but the validator should classify every form of the managed flag identically so the boundary is consistent. * Round 10: signed -np in CLI canonicaliser + reject empty inline aliases Round 10 reviewer flagged two real issues: - _expand_attached_np_short rewrote only -np; signed forms -np-1 / -np+1 fell through. Backend _flag_name already classifies them as managed, so the CLI rewriter must too -- otherwise Click clusters -np-1 into -n -p -1 (port=-1) and never reaches the backend validator at all. - -m= / -hfr= / -f= empty inline forms were accepted and produced --model '' / --frontend '' (then Path('') silently became '.') on re-exec. Reject empty inline values at the preprocessor with a clear BadParameter so the malformed input fails fast. Both behaviours pinned with parametrised regression tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Expose --parallel on plain `unsloth studio` for API-path parity The PR added --parallel to `unsloth studio run` but the plain `unsloth studio` callback (used for API-only / bare-server launches) still hardcoded llama_parallel_slots to its run_server default. With --parallel now denied as a llama_extra_args pass-through, that flow had no first-class way to raise concurrency. - unsloth_cli/commands/studio.py: add --parallel / --n-parallel typer Option (default 4, range 1..64) to studio_default, forward through the venv re-exec, and pass llama_parallel_slots= to run_server in the in-venv path. - studio/backend/run.py: argparse --parallel / --n-parallel with the same range guard so the spawned child accepts the forwarded flag. - unsloth_cli/tests/test_studio_run_parallel_flag.py: test pins the new option presence, aliases, default and range guards. * Round 12: narrow entry-point gate, preserve pre-PR plain-studio default, drop brittle source-text test Three Opus subagent reviewers (security / backcompat / code-quality) flagged the same handful of real issues. Consensus fixes: - unsloth_cli/__init__.py: narrow the -np canonicaliser gate to just {unsloth, unsloth.exe} (the only pyproject-declared console_script). The previous cli.py / unsloth-cli.py entries would silently rewrite sys.argv for any third-party myproj/cli.py that happens to import unsloth_cli. Dev users running python cli.py ... -np N still work via the space form, which parses without the rewrite. - unsloth_cli/commands/studio.py + studio/backend/run.py: restore the pre-PR llama_parallel_slots default of 1 on plain unsloth studio and python studio/backend/run.py. unsloth studio run keeps its hardcoded-pre-PR default of 4. Without this, my earlier API-path parity commit silently dropped per-call context to ctx/4 for the plain-studio flow. - unsloth_cli/tests/test_studio_run_parallel_flag.py: drop the brittle source-text grep test (test_run_kwargs_use_parallel_value). The parametrised runtime test test_in_venv_path_passes_parallel_to_run_server already pins the same intent against actual behaviour. - unsloth_cli/tests/test_studio_run_short_alias_clashes.py: pin the narrow entry-point gate with a parametrised negative test covering seven third-party argv[0] basenames (cli.py, /path/myproj/cli.py, pytest, unsloth-cli, etc.). Re-broadening the gate now trips a test instead of silently mutating an unrelated CLI's argv. * Round 13: shared parallel constants, denylist invariant test, defence-in-depth Three Opus subagent reviewers (adversarial-user / maintenance / cross-file consistency) flagged a consistent set of cleanups; folded into one commit to avoid the pre-commit.ci force-push race. unsloth_cli/commands/studio.py: - Extract _PARALLEL_MIN / _PARALLEL_MAX / _PARALLEL_DEFAULT_RUN / _PARALLEL_DEFAULT_PLAIN module-level constants and use them in both typer Options (plain studio_default = 1, studio run = 4). - _expand_attached_np_short now rewrites -np when the suffix starts with a digit (or signed digit) so '-np8x' surfaces as a clean '-np takes an int' typer error instead of a baffling '--port invalid' complaint after Click clusters '-n -p 8x'. - Re-exec forwarding emits --load-in-4bit / --no-load-in-4bit explicitly in both directions; previously the True default relied on both layers sharing the same default forever. - run() docstring now explicitly says --parallel / -np pass-through via llama_extra_args is denied (use the typer flag above). studio/backend/run.py: - Mirror the parallel constants and route the argparse default, range check, and error message through them. Help text mentions the asymmetry with 'unsloth studio run' so direct-launch dev users aren't confused by Default 1 in isolation. studio/backend/core/inference/llama_server_args.py: - _flag_name strips surrounding whitespace before denylist lookup so a caller can't slip a managed flag past the boundary with a trailing space (the trimmed form is what downstream parsers see). Tests: - New typer-aliases-subset-of-denylist invariant: every alias the typer Option claims as --parallel on run() MUST be in the backend parallel denylist group. Catches the failure mode where someone adds a new alias and forgets the boundary. - Extended denylist parametrize to cover ~14 previously untested aliases (-mu, -dr, -hfv/-hfrv/-hffv family, -mmu, full --ui group, --models-preset / --models-autoload / --no-models-autoload). - Whitespace-padded denylist rejection (' --parallel', '-np ', etc). - --load-in-4bit re-exec test pinning both polarities + default. - -np argv rewriter regression tests. - Cross-reference headers between the two test files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: repair mlx studio base export save_method (#5727) * Round 14: align backend -np recogniser with CLI rewriter + reject parent --parallel Round 14 (reviewer.py --parallel 20 with gpt-5.3-codex-spark) flagged two real P1s and a stale-rebase warning. All three addressed. - studio/backend/core/inference/llama_server_args.py: widen _flag_name so -np with trailing junk (-np8x, -np-1foo, -np+1bar, -np9zzz) classifies as managed flag -np, matching the CLI _expand_attached_np_short rewriter. Without this, POST /api/inference/load with llama_extra_args=['-np8x'] slipped past the boundary while the CLI canonicalised the same form. The two sides now agree on every digit-prefix form. - unsloth_cli/commands/studio.py: reject --parallel on the studio group when a subcommand is invoked. Pre-PR the studio callback had no --parallel; my Round 12 addition made 'unsloth studio --parallel 8 run ...' silently drop the 8 because typer doesn't propagate parent options into subcommand kwargs. Now errors with exit 2 and a message pointing the operator at the correct invocation ('unsloth studio run --parallel 8 ...'). - Picked up origin/main via merge (parent commit 0caf0526): the pre-flight stale-rebase detector found 2 lines on main in studio/backend/core/export/export.py missing from PR HEAD. Merged cleanly with no conflicts. Tests: - Parametrised denylist coverage for -np+junk forms. - New runtime test confirms exit 2 + helpful error when the group --parallel is supplied alongside an invoked subcommand. - Test that the default group --parallel value still lets a subcommand resolve (no false-positive rejection). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten code comments across --parallel PR Comment-only pass over the seven PR-touched files; trim verbose docstrings, collapse multi-line section dividers, and drop redundant prose that the code already conveys. No behaviour change. * Studio: trim remaining verbose docstrings missed in last pass Shorten the test_studio_run_parallel_flag.py module docstring and the `Re-exec arg-builder coverage` block. No behaviour change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: second comment-tightening pass across PR-touched code Trim docstrings and inline comments in studio.py, run.py, llama_server_args.py, and unsloth_cli/__init__.py. No behaviour change; all 215 tests still pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: deny --embedding / --rerank / --tools pass-through `--embedding` and `--rerank` flip llama-server into single-endpoint mode, which breaks Studio's /v1/chat/completions hop. llama-server's own `--tools` flag silently stacks on top of Studio's tool policy resolved by `--enable-tools` / `--disable-tools`. Add all three (plus the `--embeddings` / `--reranking` plural aliases) to the boundary denylist so HTTP /load and pass-through extras both reject them cleanly instead of silently desyncing the server surface. Test added to the existing `test_denylist_rejects_all_aliases` parametrize. 220 tests pass. * Studio: make PR-touched tests robust to minimal envs + Windows Two cross-OS CI findings: 1. `test_typer_parallel_aliases_are_subset_of_backend_denylist` was doing `from core.inference.llama_server_args import _DENYLIST_GROUPS` which triggers `core/inference/__init__.py` and pulls in the full backend chain (fastapi / structlog / loggers / utils.hardware). The invariant only needs the constants tuple, so load the module directly via `importlib.util.spec_from_file_location` -- the test now runs with just typer + pytest installed. 2. `test_legacy_frontend_alias_still_promotes_to_frontend` asserted the literal string `"/tmp/dist"` after the value round-trips through `Path()`. On Windows `str(Path("/tmp/dist"))` is `"\tmp\dist"`, so the assertion tripped on the same logical path. Compare via `Path(x) == Path("/tmp/dist")` so the test passes on every OS. Both surfaced by the staging-4 cross-OS CI; no production-code change. 220 tests still pass locally. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: load llama_server_args.py directly in its unit tests Same fix as the previous CLI-test commit: import the module via `importlib.util.spec_from_file_location` instead of `from core.inference.llama_server_args import ...`, so the test no longer needs the full backend chain (fastapi / structlog / loggers / utils.hardware) installed via `core/inference/__init__.py`. The boundary validator is intentionally dependency-free; its unit tests should reflect that. * Fix test_main_composer_has_dir_auto anchor after PR #5784 PR #5784 ("Improve image generation UI") rewrote the message-input textarea's static `aria-label="Message input"` into a JSX conditional `aria-label={overlay ? "Image edit instructions" : "Message input"}` but did not update the RTL bidi-attribute regression test, leaving the literal-string `find('aria-label="Message input"')` anchor with no match. The `Repo tests (CPU)` job has been red on main since. Anchor on the inner `"Message input"` string literal instead -- it survives both spellings and still pins the same textarea element so the `dir="auto"` assertion has the right block to inspect. Verified by re-running the exact CI command: 954 passed, 3 skipped, 23 deselected (was 948 passed, 1 failed). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Long Yixing --- .../core/inference/llama_server_args.py | 137 ++--- studio/backend/run.py | 24 +- .../backend/tests/test_llama_server_args.py | 163 +++++- .../test_composer_rtl_bidi_attribute.py | 6 +- unsloth_cli/__init__.py | 22 +- unsloth_cli/commands/studio.py | 281 ++++++--- .../tests/test_studio_run_parallel_flag.py | 416 +++++++++++++ .../test_studio_run_short_alias_clashes.py | 550 ++++++++++++++++++ 8 files changed, 1411 insertions(+), 188 deletions(-) create mode 100644 unsloth_cli/tests/test_studio_run_parallel_flag.py create mode 100644 unsloth_cli/tests/test_studio_run_short_alias_clashes.py diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index d8b7eb383e..4f528a689f 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -1,46 +1,29 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Validator for user-supplied llama-server pass-through args. +"""Boundary validator for user-supplied llama-server pass-through args. -Studio runs llama-server as a managed subprocess and lets callers pass -extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP: -``LoadRequest.llama_extra_args``). This module is the boundary that -rejects only flags Studio fundamentally cannot share with the user -- -model identity, the auth key, and the network endpoint Studio's HTTP -proxy targets. Anything else passes through. +Reject only flags Studio manages (model identity, auth, network, +parallel slots). Everything else (sampling, ``-c``, ``-ngl``, +``--flash-attn``, ``--cache-type-*``, ``--spec-*``, ``--jinja``, ...) +is appended after Studio's auto-set flags so llama.cpp's last-wins +parser lets the user override. -User-supplied args are appended to ``cmd`` after Studio's auto-set -flags, so llama.cpp's last-wins CLI parsing makes the user's value -override the auto-set one. That covers tunable knobs the user might -reasonably want to override -- ``-c``/``--ctx-size``, -``-np``/``--parallel``, ``-fa``/``--flash-attn``, -``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``, -``--cache-type-k/v``, ``--chat-template-file/-kwargs``, -``--spec-*``, ``--jinja``/``--no-jinja``, -``--no-context-shift``/``--context-shift``, sampling params, etc. - -Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md +Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md """ from __future__ import annotations from typing import Iterable, Optional -# Each group is the full set of aliases (short + long) for one -# hard-denied flag, taken from the llama-server README. If llama.cpp -# adds a new alias for an existing denied flag, extend the relevant -# group. -# -# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl, -# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*, -# --chat-template-*, --spec-*) pass through and override Studio's -# auto-set version via llama.cpp's last-wins CLI parsing. +# Each group = every alias (short + long) of one hard-denied flag. +# Extend the matching group when llama.cpp adds a new alias. _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( - # Model identity -- Studio resolves the model from LoadRequest and - # passes -m / mmproj after downloading from HF if needed. A second - # -m would point at a different model than the one Studio thinks - # is loaded. + # Parallel slots: owned by typer --parallel; a pass-through would + # desync app.state.llama_parallel_slots from llama-server. + frozenset({"-np", "--parallel", "--n-parallel"}), + # Model identity: Studio resolves it from LoadRequest; a second + # -m would load a different model than Studio thinks it loaded. frozenset({"-m", "--model"}), frozenset({"-mu", "--model-url"}), frozenset({"-dr", "--docker-repo"}), @@ -51,28 +34,21 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( frozenset({"-hft", "--hf-token"}), frozenset({"-mm", "--mmproj"}), frozenset({"-mmu", "--mmproj-url"}), - # Networking -- Studio binds llama-server's port and reverse-proxies - # HTTP traffic to it. Retargeting host/port/path/prefix would - # orphan Studio's proxy and the UI would lose the server. + # Networking: Studio binds + proxies; retargeting orphans the proxy. frozenset({"--host"}), frozenset({"--port"}), frozenset({"--path"}), frozenset({"--api-prefix"}), frozenset({"--reuse-port"}), - # Auth / TLS -- Studio terminates auth at its own layer; an - # upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM - # key, and TLS on llama-server would break the local proxy hop. + # Auth / TLS: Studio terminates auth; upstream --api-key / TLS + # shadows Studio's key and breaks the proxy hop. frozenset({"--api-key"}), frozenset({"--api-key-file"}), frozenset({"--ssl-key-file"}), frozenset({"--ssl-cert-file"}), - # Single-model server -- Studio runs one model per llama-server - # process and serves its own UI. Enabling multi-model loading or - # llama-server's built-in web UI changes the surface clients see. - # ``--webui``/``--no-webui`` are the legacy spelling; current - # upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions. - # Keep both so the denylist matches old and new llama-server - # binaries (Studio's prebuilt vs system-llama.cpp). + # Built-in web UI. --webui/--no-webui is the legacy spelling; + # upstream renamed to --ui/--no-ui + --ui-*. Keep both so prebuilt + # and system llama.cpp binaries both match. frozenset({"--webui", "--no-webui"}), frozenset({"--ui", "--no-ui"}), frozenset({"--ui-config"}), @@ -82,32 +58,46 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( frozenset({"--models-preset"}), frozenset({"--models-max"}), frozenset({"--models-autoload", "--no-models-autoload"}), + # Server-mode flips: --embedding / --rerank restrict llama-server to + # those endpoints, breaking Studio's /v1/chat/completions hop. + frozenset({"--embedding", "--embeddings"}), + frozenset({"--rerank", "--reranking"}), + # llama-server's own built-in tools flag would silently stack on top + # of Studio's --enable-tools / --disable-tools policy resolver. + frozenset({"--tools"}), ) _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS) def _flag_name(token: str) -> Optional[str]: - """Return the flag name for a token, or None if it isn't a flag. + """Flag name for ``token``, or None if it isn't a flag. - Peels ``--key=value`` to the bare ``--key``. Plain numeric values - like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags; - llama-server short-form flags always start with a letter. + Peels `--key=value` to `--key`, treats `-1` / `-0.5` as values + (llama-server shorts always start with a letter), strips + whitespace, and normalises attached `-np8` / signed `-np-1` / + digit-prefix-junk `-np8x` to `-np`. Mirrors the CLI's + `_expand_attached_np_short`. """ + token = token.strip() if not token.startswith("-") or token in {"-", "--"}: return None if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): return None - return token.split("=", 1)[0] + name = token.split("=", 1)[0] + if len(name) > 3 and name.startswith("-np"): + suffix = name[3:] + if suffix[0].isdigit() or ( + len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit() + ): + return "-np" + return name def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: - """Validate user-supplied llama-server args. - - Returns the args as a flat list ready to extend the llama-server - command. Raises ``ValueError`` (with the offending flag in the - message) the moment a token resolves to a Studio-managed flag. - """ + """Validate user-supplied llama-server args. Returns a flat list + ready to extend the llama-server command; raises ``ValueError`` + naming the offending flag on the first managed token.""" if not args: return [] out: list[str] = [] @@ -124,15 +114,15 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: def is_managed_flag(flag: str) -> bool: - """True if ``flag`` is a Studio-managed llama-server flag.""" - return flag in _DENYLIST + """True if ``flag`` is Studio-managed. Normalises via ``_flag_name`` + so `-np8` / `--parallel=8` classify like the canonical tokens.""" + normalised = _flag_name(flag) + return normalised is not None and normalised in _DENYLIST -# Pass-through flags that shadow first-class ``LoadRequest`` fields -# (max_seq_length, cache_type_kv, speculative_type, -# chat_template_override). Stripped from inherited extras so they -# can't last-wins-override an Apply that re-sets the same first-class -# field. +# Pass-through flags that shadow first-class LoadRequest fields; +# stripped from inherited extras so they can't last-wins-override an +# Apply that re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) _CACHE_FLAGS: frozenset[str] = frozenset( {"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"} @@ -169,9 +159,8 @@ _SHADOWING_FLAGS: frozenset[str] = ( _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS ) -# Boolean flags inside _SHADOWING_FLAGS that take no value. The -# value-consuming heuristic in strip_shadowing_flags must skip just the -# flag for these, never the following token. +# Shadowing flags that take no value -- strip the flag only, never the +# following token. _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( {"--spec-default", "--jinja", "--no-jinja"} ) @@ -187,14 +176,11 @@ def strip_shadowing_flags( ) -> list[str]: """Strip flags that shadow first-class Studio settings. - Used when the route inherits a previous load's ``llama_extra_args`` - so that an inherited ``-c 4096`` cannot override the current - request's ``max_seq_length`` (and equivalents for cache / - speculative / chat template). Each ``strip_*`` flag controls one - group; the route only strips groups whose corresponding first-class - field was actually supplied by the caller, so an inherited - ``--chat-template-file`` survives an Apply that omits both - ``llama_extra_args`` and ``chat_template_override``. + Used when inheriting a previous load's ``llama_extra_args`` so an + inherited `-c 4096` can't override the current `max_seq_length` + (same for cache / spec / template). Each ``strip_*`` toggle + controls one group; the route only strips groups whose first-class + field the caller actually supplied. """ shadowing: set[str] = set() if strip_context: @@ -216,9 +202,8 @@ def strip_shadowing_flags( out.append(tok) i += 1 continue - # Drop this token. Boolean shadowing flags never carry a value; - # other shadowing flags consume the next token when it isn't a - # flag and the value isn't already packed as ``--key=value``. + # Drop the flag; consume the next token too unless it's + # boolean, already inline (`-c=4096`), or another flag. if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok: i += 1 elif i + 1 < n and _flag_name(tokens[i + 1]) is None: diff --git a/studio/backend/run.py b/studio/backend/run.py index 3bde8abd3c..e96e609659 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -846,11 +846,33 @@ if __name__ == "__main__": action = "store_true", help = "API server only, no frontend (for Tauri)", ) + # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 + # applies only to direct backend launches; `unsloth studio run` + # always passes its own value (4) explicitly. + _PARALLEL_MIN = 1 + _PARALLEL_MAX = 64 + _PARALLEL_DEFAULT_PLAIN = 1 + parser.add_argument( + "--parallel", + "--n-parallel", + type = int, + default = _PARALLEL_DEFAULT_PLAIN, + help = ( + f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). " + f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4." + ), + ) args = parser.parse_args() + if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX: + parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}") kwargs = dict( - host = args.host, port = args.port, silent = args.silent, api_only = args.api_only + host = args.host, + port = args.port, + silent = args.silent, + api_only = args.api_only, + llama_parallel_slots = args.parallel, ) if args.frontend is not None: kwargs["frontend_path"] = Path(args.frontend) diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 68a1c870fb..02a272ba3e 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -3,21 +3,35 @@ """Unit tests for the llama-server pass-through args validator. -The validator is the security boundary between user-supplied CLI / HTTP -input and the llama-server subprocess command. These tests pin the -denylist behavior so the boundary doesn't quietly regress when new -managed flags are added. +The validator is the boundary between user CLI/HTTP input and the +llama-server subprocess. These tests pin denylist behaviour so it +doesn't quietly regress when new managed flags are added. """ from __future__ import annotations +import importlib.util +import re +from pathlib import Path + import pytest -from core.inference.llama_server_args import ( - is_managed_flag, - strip_shadowing_flags, - validate_extra_args, +# Load llama_server_args.py directly so this test doesn't drag in the +# full backend chain (fastapi / structlog / loggers / utils.hardware) +# via core/inference/__init__.py. The validator is intentionally +# dependency-free and unit-tests should reflect that. +_LSA_PATH = ( + Path(__file__).resolve().parent.parent + / "core" + / "inference" + / "llama_server_args.py" ) +_spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH) +_lsa = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_lsa) +is_managed_flag = _lsa.is_managed_flag +strip_shadowing_flags = _lsa.strip_shadowing_flags +validate_extra_args = _lsa.validate_extra_args # ── Pass-through (allowed) ─────────────────────────────────────────── @@ -60,13 +74,12 @@ from core.inference.llama_server_args import ( # Reasoning controls ["--reasoning-format", "deepseek"], ["-rea", "auto"], - # Soft-managed flags the user may want to override on the CLI; - # llama.cpp's last-wins parsing means these win over Studio's - # auto-set version. + # Soft-managed: user-supplied flags last-wins-override Studio's + # auto-set version. --parallel / -np / --n-parallel are NOT + # here -- they're hard-denied (KV-cache + slot count would + # desync). Use `unsloth studio run --parallel N` instead. ["-c", "131072"], ["--ctx-size", "8192"], - ["--parallel", "1"], - ["-np", "8"], ["--flash-attn", "off"], ["-fa", "on"], ["--no-context-shift"], @@ -99,8 +112,7 @@ def test_value_with_equals_form_passes_through(): def test_non_flag_token_passes_through(): - # A bare positional value (not preceded by a flag) is preserved - # verbatim. llama-server may reject it, but that's not our job. + # Bare positionals are passed through; llama-server can reject them. assert validate_extra_args(["foo"]) == ["foo"] @@ -110,18 +122,33 @@ def test_non_flag_token_passes_through(): @pytest.mark.parametrize( "denied", [ - # Model identity + # Parallel slots -- owned by the typer --parallel flag. + "-np", + "--parallel", + "--n-parallel", + # Model identity (every alias; bumping llama.cpp must keep + # every form rejected, not just the long). "-m", "--model", + "-mu", + "--model-url", + "-dr", + "--docker-repo", "-hf", "-hfr", "--hf-repo", "-hff", "--hf-file", + "-hfv", + "-hfrv", + "--hf-repo-v", + "-hffv", + "--hf-file-v", "-hft", "--hf-token", "-mm", "--mmproj", + "-mmu", "--mmproj-url", # Networking (Studio binds + proxies) "--host", @@ -134,11 +161,28 @@ def test_non_flag_token_passes_through(): "--api-key-file", "--ssl-key-file", "--ssl-cert-file", - # Single-model server + # Single-model server (legacy --webui + current --ui group) "--webui", "--no-webui", + "--ui", + "--no-ui", + "--ui-config", + "--ui-config-file", + "--ui-mcp-proxy", + "--no-ui-mcp-proxy", "--models-dir", + "--models-preset", "--models-max", + "--models-autoload", + "--no-models-autoload", + # Server-mode flips: --embedding / --rerank would restrict + # llama-server to those endpoints and break Studio's chat hop. + "--embedding", + "--embeddings", + "--rerank", + "--reranking", + # llama-server's own --tools clashes with Studio's tool policy. + "--tools", ], ) def test_denylist_rejects_all_aliases(denied): @@ -146,14 +190,65 @@ def test_denylist_rejects_all_aliases(denied): validate_extra_args([denied, "value"]) +@pytest.mark.parametrize( + "args,offending", + [ + # Pass-through --parallel would last-wins-override the real + # slot count while Studio's KV-cache fit + llama_parallel_slots + # stay at the typer value -- plan vs. process disagree. + (["--parallel", "8"], "--parallel"), + (["--parallel=8"], "--parallel"), + (["--n-parallel", "16"], "--n-parallel"), + (["--n-parallel=16"], "--n-parallel"), + (["-np", "32"], "-np"), + # Attached short form: Click clusters it CLI-side; HTTP /load + # with `["-np8"]` must still resolve to managed. + (["-np8"], "-np"), + (["-np64"], "-np"), + # Out-of-range values that would bypass the typer 1..64 guard. + (["--parallel", "999"], "--parallel"), + (["-np", "0"], "-np"), + (["-np999"], "-np"), + # Signed attached forms; `-np-1` must not slip past. + (["-np-1"], "-np"), + (["-np+1"], "-np"), + ], +) +def test_parallel_flags_are_managed(args, offending): + with pytest.raises(ValueError, match = re.escape(offending)): + validate_extra_args(args) + + def test_denylist_rejects_equals_form(): with pytest.raises(ValueError, match = "--port"): validate_extra_args(["--port=9000"]) +@pytest.mark.parametrize( + "padded", + [" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"], +) +def test_denylist_rejects_whitespace_padded_forms(padded): + # `_flag_name` trims whitespace before lookup; otherwise a trailing + # space could slip a managed flag past the boundary. + with pytest.raises(ValueError, match = "parallel|np"): + validate_extra_args([padded, "8"]) + + +@pytest.mark.parametrize( + "attached", + ["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"], +) +def test_denylist_rejects_np_with_digit_prefix_and_junk(attached): + # Backend `_flag_name` must classify the same forms the CLI + # rewriter expands, else HTTP /load could smuggle `-np8x` through. + with pytest.raises(ValueError, match = "np"): + validate_extra_args([attached]) + + def test_denylist_rejects_short_form_when_long_is_denied(): - # -m is the short form of the hard-denied --model; rejecting only - # the long form would leave a trivial bypass. + # `-m` is the short form of --model; rejecting only the long + # form would leave a trivial bypass. with pytest.raises(ValueError, match = "-m"): validate_extra_args(["-m", "/some/other/path.gguf"]) @@ -165,9 +260,7 @@ def test_denylist_message_names_offending_flag(): def test_first_denied_flag_short_circuits(): - # Validation stops at the first denied flag; later denied flags - # in the same call don't matter for behaviour, but the message - # should name the first one we hit. + # Validation stops at the first denied flag; the message names it. with pytest.raises(ValueError, match = "--port"): validate_extra_args(["--port", "1", "--host", "x"]) @@ -177,8 +270,7 @@ def test_first_denied_flag_short_circuits(): @pytest.mark.parametrize("value", ["-1", "-0.5", "-42", "-.5"]) def test_negative_number_value_is_not_flag(value): - # ``--seed -1`` is a value, not a flag. Validator must not try - # to look up "-1" in the denylist. + # `--seed -1`: the -1 is a value, not a flag. assert validate_extra_args(["--seed", value]) == ["--seed", value] @@ -190,6 +282,15 @@ def test_is_managed_flag_true_for_denied(): assert is_managed_flag("--api-key") is True assert is_managed_flag("-m") is True assert is_managed_flag("--model") is True + # Parallel slots owned by the typer --parallel flag. + assert is_managed_flag("--parallel") is True + assert is_managed_flag("--n-parallel") is True + assert is_managed_flag("-np") is True + # Normalised forms must classify like the canonical token so + # is_managed_flag filtering stays in sync with validate_extra_args. + assert is_managed_flag("-np8") is True + assert is_managed_flag("--parallel=8") is True + assert is_managed_flag("--port=9000") is True def test_is_managed_flag_false_for_pass_through(): @@ -199,7 +300,6 @@ def test_is_managed_flag_false_for_pass_through(): # Soft-managed flags pass through (last-wins override) assert is_managed_flag("-c") is False assert is_managed_flag("--ctx-size") is False - assert is_managed_flag("--parallel") is False assert is_managed_flag("--flash-attn") is False assert is_managed_flag("-ngl") is False assert is_managed_flag("--threads") is False @@ -231,8 +331,8 @@ def test_strip_shadowing_flags_keeps_context_when_not_requested(): def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled(): - # Caller did not supply chat_template_override; the inherited - # --chat-template-file must survive the strip. + # No chat_template_override supplied; inherited + # --chat-template-file must survive. out = strip_shadowing_flags( ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"], strip_context = True, @@ -282,7 +382,7 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled(): def test_strip_shadowing_flags_drops_mtp_flags_when_requested(): - # MTP / draft-mtp flags must be stripped when speculative_type is re-applied. + # MTP / draft-mtp flags must drop when speculative_type re-applies. out = strip_shadowing_flags( [ "--spec-type", @@ -311,8 +411,7 @@ def test_is_managed_flag_false_for_mtp_pass_through(): def test_strip_shadowing_flags_boolean_does_not_consume_next_token(): - # --spec-default is a boolean shadowing flag; the value-skipping - # heuristic must skip just the flag, not the following positional. + # `--spec-default` is boolean; drop just the flag, keep the next token. out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True) assert out == ["ngram-mod"] @@ -343,8 +442,8 @@ def test_strip_shadowing_flags_handles_empty_input(): def test_strip_shadowing_flags_defaults_strip_everything(): - # The route's already-loaded comparator calls strip_shadowing_flags - # with no kwargs to detect ANY shadowing flag in stored extras. + # The route's already-loaded comparator calls with no kwargs to + # detect ANY shadowing flag in stored extras. out = strip_shadowing_flags( ["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"] ) diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py index a1af16d4fc..1e2f0d0cf4 100644 --- a/tests/studio/test_composer_rtl_bidi_attribute.py +++ b/tests/studio/test_composer_rtl_bidi_attribute.py @@ -26,7 +26,11 @@ def _block_around(src: str, anchor: str, radius: int = 600) -> str: def test_main_composer_has_dir_auto(): - block = _block_around(THREAD_TSX.read_text(), 'aria-label="Message input"') + # PR #5784 rewrote the literal attribute into a JSX conditional + # (`aria-label={overlay ? "Image edit instructions" : "Message input"}`), + # so anchor on the inner string literal instead -- it survives both + # the old and new spellings. + block = _block_around(THREAD_TSX.read_text(), '"Message input"') assert 'dir="auto"' in block, 'main composer is missing dir="auto"' diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 65834b3b9d..c8ec4c66c0 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import os.path as _osp +import sys as _sys + import typer from importlib.metadata import version as package_version, PackageNotFoundError @@ -8,7 +11,19 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference from unsloth_cli.commands.export import export, list_checkpoints -from unsloth_cli.commands.studio import run as studio_run, studio_app +from unsloth_cli.commands.studio import ( + run as studio_run, + studio_app, + _expand_attached_np_short, +) + + +# Canonicalise `-np` only under the `unsloth` console-script; +# third-party scripts that import unsloth_cli keep their argv intact. +_entry_base = _osp.basename(_sys.argv[0]).lower() if _sys.argv else "" +if _entry_base in {"unsloth", "unsloth.exe"}: + _expand_attached_np_short() +del _entry_base def show_version(value: bool): @@ -47,9 +62,8 @@ app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") -# Top-level alias: `unsloth run ...` is equivalent to `unsloth studio run ...`. -# Same context_settings as the studio_app registration so unknown flags -# still pass through to llama-server. +# Top-level `unsloth run` aliases `unsloth studio run`; same context +# so unknown flags still pass through to llama-server. app.command( "run", context_settings = { diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index e37cd0a8d8..edb8ea2cf4 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -206,6 +206,14 @@ def _find_setup_script() -> Optional[Path]: return None +# Mirror in studio/backend/run.py argparse + backend denylist test; +# bumping the cap in one place only desyncs. +_PARALLEL_MIN = 1 +_PARALLEL_MAX = 64 +_PARALLEL_DEFAULT_RUN = 4 # pre-PR hardcoded for `unsloth studio run` +_PARALLEL_DEFAULT_PLAIN = 1 # pre-PR effective for plain `unsloth studio` + + def _iter_editable_studio_source_roots(venv_dir: Path): """Yield repo roots from setuptools `__editable___*_finder.py` files in *venv_dir*'s site-packages whose MAPPING includes a `studio` entry. @@ -587,14 +595,38 @@ def studio_default( "--api-only", help = "Run API server only, no frontend serving (for Tauri desktop app)", ), + parallel: int = typer.Option( + _PARALLEL_DEFAULT_PLAIN, + "--parallel", + "--n-parallel", + min = _PARALLEL_MIN, + max = _PARALLEL_MAX, + help = ( + f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). " + f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` " + f"defaults to {_PARALLEL_DEFAULT_RUN}." + ), + ), ): """Launch the Unsloth Studio server.""" - # Runs before any subcommand; covers run/setup/update/etc in one place. + # Runs before every subcommand (run/setup/update/...). _ensure_studio_env_exported() if ctx.invoked_subcommand is not None: + # Typer doesn't forward parent options to subcommands, so + # `unsloth studio --parallel N run ...` would silently drop N. + if parallel != _PARALLEL_DEFAULT_PLAIN: + typer.echo( + f"Error: --parallel on `unsloth studio` applies to the " + f"plain-server path only. For `unsloth studio " + f"{ctx.invoked_subcommand}`, put the flag after the " + f"subcommand: `unsloth studio {ctx.invoked_subcommand} " + f"--parallel {parallel} ...`", + err = True, + ) + raise typer.Exit(2) return - # Always use the studio venv if it exists and we're not already in it + # Use the studio venv if it exists and we aren't already in it. studio_venv_dir = STUDIO_HOME / "unsloth_studio" in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) @@ -611,6 +643,8 @@ def studio_default( host, "--port", str(port), + "--parallel", + str(parallel), ] # Resolve frontend explicitly so the spawned run.py uses a real # built dist regardless of where its __file__ lands. Skip in @@ -624,9 +658,8 @@ def studio_default( args.append("--silent") if api_only: args.append("--api-only") - # On Windows, os.execvp() spawns a child but the parent lingers, - # so Ctrl+C only kills the parent leaving the child orphaned. - # Use subprocess.run() on Windows so the parent waits for the child. + # On Windows os.execvp keeps the parent alive, so Ctrl+C + # would orphan the child; use Popen+wait instead. if sys.platform == "win32": import subprocess as _sp @@ -634,7 +667,7 @@ def studio_default( try: rc = proc.wait() except KeyboardInterrupt: - # Child has its own signal handler — let it finish + # Child handles its own signal; let it finish. rc = proc.wait() if rc != 0: typer.echo( @@ -661,7 +694,13 @@ def studio_default( display_host = _resolve_external_ip() if host == "0.0.0.0" else host typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}") - run_kwargs = dict(host = host, port = port, silent = silent, api_only = api_only) + run_kwargs = dict( + host = host, + port = port, + silent = silent, + api_only = api_only, + llama_parallel_slots = parallel, + ) if frontend is not None: run_kwargs["frontend_path"] = frontend run_server(**run_kwargs) @@ -670,8 +709,8 @@ def studio_default( try: if _shutdown_event is not None: - # NOTE: Event.wait() without a timeout blocks at the C level - # on Linux, preventing Python from delivering SIGINT (Ctrl+C). + # Event.wait() with no timeout blocks at C-level on Linux + # and swallows SIGINT; loop with a 1s timeout instead. while not _shutdown_event.is_set(): _shutdown_event.wait(timeout = 1) else: @@ -688,21 +727,15 @@ def studio_default( def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]: - """Split ``org/name:variant`` HF-style identifiers into (repo, variant). - - Mirrors llama.cpp's ``-hf :`` convention so users can - write ``unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL`` instead of passing - ``--gguf-variant`` separately. Local paths (absolute, ``./``, - ``~/``, Windows drive letters) and identifiers without a ``:`` - suffix are returned verbatim. - """ + """Split ``org/name:variant`` into ``(repo, variant)``; mirrors + llama.cpp's ``-hf :``. Local paths, Windows drives, + and ids without ``:`` pass through verbatim.""" s = model_arg.strip() if not s: return s, None if s.startswith(("/", "./", "../", "~")) or s == ".": return s, None - # Windows drive letter (e.g. "C:\\path" or "C:/path") -- the colon - # here is a path separator, not a variant suffix. + # Windows drive letter (e.g. "C:\path"): colon is a path separator. if len(s) >= 2 and s[1] == ":" and s[0].isalpha(): return s, None if ":" not in s: @@ -710,13 +743,80 @@ def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]: repo, _, variant = s.rpartition(":") if not repo or not variant: return s, None - # A real quant label has no slashes; ``foo:bar/baz`` is not - # ``repo:variant`` syntax. + # Quant labels never contain a slash; `foo:bar/baz` isn't repo:variant. if "/" in variant: return s, None return repo, variant +def _expand_attached_np_short() -> None: + # Click clusters `-np8` as `-n -p 8` (-p = --port), dropping the + # parallel value. Split to `-np ` so typer's alias matches. + # Stops at `--`; accepts signed and digit-prefix-junk forms so + # typer can report a clean error against `-np`. Kept in lockstep + # with the backend `_flag_name` recogniser. + i = 0 + while i < len(sys.argv): + tok = sys.argv[i] + if tok == "--": + break + if len(tok) > 3 and tok.startswith("-np") and tok[3] != "=": + suffix = tok[3:] + first_numeric = suffix[0].isdigit() or ( + len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit() + ) + if first_numeric: + sys.argv[i : i + 1] = ["-np", suffix] + i += 2 + continue + i += 1 + + +def _consume_legacy_short_aliases( + args: List[str], + aliases: tuple[str, ...], + current: Optional[str], + canonical: str, +) -> tuple[Optional[str], List[str]]: + """Pop exact-match legacy shorts (`-m`/`-hfr`/`-f`) from args; + leave clusters (`-mg`/`-fa`/...) for the llama-server tail. Inline + `-x=value` form also accepted.""" + out: List[str] = [] + value = current + i, n = 0, len(args) + while i < n: + tok = args[i] + if tok == "--": # end of options; tail is raw payload. + out.extend(args[i:]) + break + name, sep, inline = tok.partition("=") + if name not in aliases: + out.append(tok) + i += 1 + continue + if value is not None: + raise typer.BadParameter( + f"{name} conflicts with {canonical} already provided" + ) + if sep: + if inline == "": # `-m=` would become --model '' (Path('')='.'). + raise typer.BadParameter(f"{name} requires a non-empty value") + value = inline + i += 1 + elif i + 1 < n: + nxt = args[i + 1] + # `--long` is unambiguously a flag; single-dash `-x` may be a path. + if nxt.startswith("--") and nxt != "--": + raise typer.BadParameter( + f"{name} expects a value but got the flag {nxt}" + ) + value = nxt + i += 2 + else: + raise typer.BadParameter(f"{name} requires a value") + return value, out + + @studio_app.command( context_settings = { "allow_extra_args": True, @@ -725,17 +825,18 @@ def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]: ) def run( ctx: typer.Context, - model: str = typer.Option( - ..., + model: Optional[str] = typer.Option( + None, "--model", - "-m", "-hf", - "-hfr", "--hf-repo", + # `-m` / `-hfr` removed (Click would cluster `-mg`/`-md`/...). + # Exact-match `-m`/`-hfr` still work via the legacy shim below. + # `-hf` stays (multi-char shorts don't cluster). help = ( "Model path or HF repo. Accepts llama.cpp-style " - "`org/repo:variant` syntax. The `-hf` / `--hf-repo` aliases " - "match llama-server's spelling." + "`org/repo:variant` syntax. `-hf` / `--hf-repo` match " + "llama-server's spelling." ), ), gguf_variant: Optional[str] = typer.Option( @@ -750,7 +851,8 @@ def run( ), port: int = typer.Option(8888, "--port", "-p"), host: str = typer.Option("127.0.0.1", "--host", "-H"), - frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"), + # `-f` removed (clustered `-fa`/`-fit*`); studio_default keeps it. + frontend: Optional[Path] = typer.Option(None, "--frontend"), silent: bool = typer.Option(False, "--silent", "-q"), enable_tools: Optional[bool] = typer.Option( None, @@ -766,26 +868,65 @@ def run( "-y", help = "Skip the 0.0.0.0 + --enable-tools confirmation prompt.", ), + parallel: int = typer.Option( + _PARALLEL_DEFAULT_RUN, + "--parallel", + "--n-parallel", + "-np", + min = _PARALLEL_MIN, + max = _PARALLEL_MAX, + help = ( + "llama-server parallel decode slots. N requests share one " + "loaded model; each slot gets ctx/N KV cache. Default " + f"{_PARALLEL_DEFAULT_RUN} (pre-PR hardcoded value)." + ), + ), ): - """Start Studio, load a model, and print an API key -- one-liner server. + """Start Studio, load a model, print an API key -- one-liner server. - Any flag this command does not recognize is forwarded verbatim to - the underlying llama-server (GGUF only). Studio-managed flags - (--port, -c / --ctx-size, --api-key, -ngl, --jinja, --flash-attn, - --no-context-shift, model-identity flags, ...) are rejected with - HTTP 400. + Unknown flags pass through to llama-server (GGUF only). Studio + rejects managed flags with HTTP 400: model identity, network + (--host/--port/--path/--api-prefix/--reuse-port), auth/TLS + (--api-key/--ssl-*), single-model UI (--ui/--models-*/--webui), + and parallel slots (use --parallel above). Full denylist in + studio/backend/core/inference/llama_server_args.py. Other knobs + (-c, -ngl, --jinja, --flash-attn, -t, ...) pass through and + last-wins-override Studio's auto-set value. Example: unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL - unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42 + unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42 --parallel 8 unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja """ extra_llama_args: List[str] = list(ctx.args) if ctx.args else [] - # ── 0. Parse llama.cpp-style ``repo:variant`` syntax in --model. ─── - # Lets users write ``--model unsloth/foo-GGUF:UD-Q4_K_XL`` instead - # of pairing ``--model`` with ``--gguf-variant``. If both are given - # and disagree, fail loudly instead of silently picking one. + # Promote legacy exact `-m`/`-hfr`/`-f` back into typer params; + # clusters stay in extras. + model, extra_llama_args = _consume_legacy_short_aliases( + extra_llama_args, + ("-m", "-hfr"), + model, + "--model", + ) + legacy_frontend, extra_llama_args = _consume_legacy_short_aliases( + extra_llama_args, + ("-f",), + str(frontend) if frontend is not None else None, + "--frontend", + ) + if legacy_frontend is not None and frontend is None: + frontend = Path(legacy_frontend) + + if model is None: + typer.echo( + "Error: Missing option '--model' / '-hf' / '--hf-repo' " + "(legacy aliases '-m' / '-hfr' are still accepted).", + err = True, + ) + raise typer.Exit(2) + + # 0. Parse llama.cpp `repo:variant` in --model; error if also paired + # with --gguf-variant and they disagree. parsed_repo, embedded_variant = _split_repo_variant(model) if embedded_variant: if gguf_variant and gguf_variant != embedded_variant: @@ -798,8 +939,8 @@ def run( model = parsed_repo gguf_variant = gguf_variant or embedded_variant - # ── Resolve the server-side tool policy. The y/N prompt (if any) - # runs in the outer process so the re-exec'd child never re-prompts. + # Resolve tool policy here so the re-exec'd child inherits a + # concrete decision and never re-prompts. from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy enable_tools = resolve_tool_policy( @@ -809,7 +950,7 @@ def run( silent = silent, ) - # ── 1. Venv re-exec (same pattern as studio_default) ────────────── + # 1. Re-exec into the studio venv (same pattern as studio_default). studio_venv_dir = STUDIO_HOME / "unsloth_studio" in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) @@ -818,7 +959,7 @@ def run( if not studio_python: typer.echo("Studio not set up. Run install.sh first.") raise typer.Exit(1) - # Re-exec into the studio venv via its `unsloth` entry point + # Re-exec via the studio venv's `unsloth` console-script. studio_bin = studio_python.parent / "unsloth" if not studio_bin.is_file(): typer.echo( @@ -842,27 +983,26 @@ def run( ] if gguf_variant: args.extend(["--gguf-variant", gguf_variant]) - if not load_in_4bit: - args.append("--no-load-in-4bit") + # Forward the explicit polarity; a future default flip on one + # layer must not silently invert behaviour for the other. + args.append("--load-in-4bit" if load_in_4bit else "--no-load-in-4bit") if frontend: args.extend(["--frontend", str(frontend)]) if silent: args.append("--silent") - # Forward the resolved tool policy (always concrete True/False - # at this point — the resolver above ran before the re-exec). + # Forward the resolved tool policy so the child doesn't re-resolve. if enable_tools: args.append("--enable-tools") else: args.append("--disable-tools") - # Forward --yes whenever the parent already cleared the prompt - # (either operator passed --yes, or the parent's resolver - # accepted the network-bind confirmation). Otherwise the child - # re-runs the resolver and prompts a second time. + # Forward --yes if the parent already cleared the network-bind + # prompt, else the child re-prompts. if yes or (enable_tools and is_external_host(host)): args.append("--yes") - # Forward unknown args (llama-server pass-through) to the - # re-exec'd command so the studio venv sees them in ctx.args - # and the re-execed run() can include them in the load payload. + # Typer claims --parallel outside ctx.args; without this the + # child reverts to its default and silently drops the value. + args.extend(["--parallel", str(parallel)]) + # llama-server pass-through extras → child ctx.args → load payload. if extra_llama_args: args.extend(extra_llama_args) @@ -879,34 +1019,31 @@ def run( # ── 2. Start server (always suppress built-in banner) ───────────── from studio.backend.run import run_server, _resolve_external_ip - run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = 4) + run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = parallel) if frontend is not None: run_kwargs["frontend_path"] = frontend app = run_server(**run_kwargs) actual_port = getattr(app.state, "server_port", port) or port - # ── Apply the resolved tool policy as a process-level override. - # Must use the same import path the route handlers use -- - # `studio/backend/run.py` adds `studio/backend/` to sys.path so the - # routes import this module as top-level `state.tool_policy`. If we - # imported via `studio.backend.state.tool_policy` instead, Python - # would cache two different module objects with two different - # `_tool_policy` globals, and the gates would never see our value. + # Match the route handlers' import path: run.py adds + # studio/backend/ to sys.path, so they import as `state.tool_policy`. + # Importing via `studio.backend.state.tool_policy` would cache a + # second module object whose flag the gates can't see. from state.tool_policy import set_tool_policy set_tool_policy(enable_tools) - # ── 3. Wait for server health ───────────────────────────────────── + # 3. Wait for server health. if not silent: typer.echo("Starting Unsloth Studio...") if not _wait_for_server(actual_port): typer.echo("Error: server did not become healthy within 30 seconds.", err = True) raise typer.Exit(1) - # ── 4. Create API key in-process ────────────────────────────────── + # 4. Create API key in-process. api_key = _create_api_key_inprocess(api_key_name) - # ── 5. Load model via HTTP ──────────────────────────────────────── + # 5. Load model via HTTP. if not silent: typer.echo(f"Loading model: {model}...") try: @@ -926,15 +1063,13 @@ def run( loaded_model = result.get("model", model) display_variant = f" ({gguf_variant})" if gguf_variant else "" - # ── 6. Print banner ─────────────────────────────────────────────── + # 6. Print banner. display_host = _resolve_external_ip() if host == "0.0.0.0" else host base_url = f"http://{display_host}:{actual_port}" sdk_base_url = f"{base_url}/v1" - # Claude orange (Claude Code's brand color) for tool-policy notices - # so they stand out from the surrounding banner. Always printed -- - # even under --silent / --yes -- so the operator never misses the - # current tool-execution status. + # Orange so the tool-policy notice stands out; printed under + # --silent / --yes too so the policy is never invisible. _tool_notice_fg = (217, 119, 87) _is_external = is_external_host(host) if _is_external and enable_tools: @@ -992,14 +1127,12 @@ def run( typer.echo(""" -d '{"input": "Hello", "stream": true}'""") typer.echo("") else: - # Silent mode still prints the essentials (URL, API key) plus - # the orange tool-status notice so the operator never loses - # visibility into the security-relevant policy. + # Silent still prints URL + API key + tool-status policy. typer.echo(f"URL: {base_url}") typer.echo(f"API Key: {api_key}") typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True) - # ── 7. Wait for Ctrl+C ──────────────────────────────────────────── + # 7. Wait for Ctrl+C. from studio.backend.run import _shutdown_event, _graceful_shutdown, _server try: diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py new file mode 100644 index 0000000000..561600e43b --- /dev/null +++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py @@ -0,0 +1,416 @@ +# 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 `unsloth studio run --parallel` CLI flag. + +Pre-PR `llama_parallel_slots` was hardcoded to 4. These tests pin +the typer Option (aliases, default 4, 1..64 range), the +typer/denylist subset invariant, and re-exec forwarding. + +See ``test_studio_run_short_alias_clashes.py`` for the argv +canonicaliser and the legacy `-m` / `-hfr` / `-f` shim. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from typer.testing import CliRunner + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _load_run_command(): + """Import `studio` without triggering server start; backend imports + are lazy inside run().""" + from unsloth_cli.commands import studio as _studio + + return _studio + + +def test_parallel_option_is_registered(): + """The `--parallel` flag (with aliases) must be on the `run` command.""" + studio_mod = _load_run_command() + import inspect + + run_fn = studio_mod.run + sig = inspect.signature(run_fn) + assert "parallel" in sig.parameters, "missing `parallel` parameter on run()" + + param = sig.parameters["parallel"] + opt = param.default # typer.OptionInfo + flags = set() + decls = getattr(opt, "param_decls", None) or [] + for d in decls: + flags.add(d) + for required in ("--parallel", "--n-parallel", "-np"): + assert required in flags, f"flag {required!r} missing from --parallel option" + + +def test_parallel_default_is_four(): + """Default must stay at 4 so plain `unsloth studio run` is unchanged.""" + studio_mod = _load_run_command() + import inspect + + sig = inspect.signature(studio_mod.run) + opt = sig.parameters["parallel"].default + default = getattr(opt, "default", None) + assert ( + default == 4 + ), f"default changed to {default}; would silently alter existing deployments" + + +def test_parallel_range_guards_are_set(): + """Range guards: 1 <= N <= 64. Outside this is a hard reject.""" + studio_mod = _load_run_command() + import inspect + + sig = inspect.signature(studio_mod.run) + opt = sig.parameters["parallel"].default + assert getattr(opt, "min", None) == 1, "min must be 1 (0 = no decode possible)" + assert getattr(opt, "max", None) == 64, "max must be 64 (KV split sanity cap)" + + +def test_typer_parallel_aliases_are_subset_of_backend_denylist(): + """Every typer alias for --parallel must be denied on the backend + too; otherwise HTTP /load could smuggle the value via + `llama_extra_args` and desync llama_parallel_slots from the + running llama-server.""" + studio_mod = _load_run_command() + import inspect + import importlib.util + + # Load llama_server_args.py directly so the test doesn't need the + # backend's full runtime chain (fastapi / structlog / loggers / + # utils.hardware) installed -- the invariant is just about the + # _DENYLIST_GROUPS tuple. + lsa_path = ( + Path(__file__).resolve().parents[2] + / "studio" + / "backend" + / "core" + / "inference" + / "llama_server_args.py" + ) + spec = importlib.util.spec_from_file_location("_lsa_for_subset_test", lsa_path) + lsa = importlib.util.module_from_spec(spec) + spec.loader.exec_module(lsa) + _DENYLIST_GROUPS = lsa._DENYLIST_GROUPS + + parallel_group = next((g for g in _DENYLIST_GROUPS if "--parallel" in g), None) + assert parallel_group is not None, "denylist must include a --parallel group" + + sig = inspect.signature(studio_mod.run) + opt = sig.parameters["parallel"].default + typer_aliases = set(getattr(opt, "param_decls", []) or []) + missing = typer_aliases - parallel_group + assert not missing, ( + f"typer aliases {missing!r} are not in the backend denylist; " + f"add them to _DENYLIST_GROUPS to keep /load from desyncing " + f"llama_parallel_slots." + ) + + +# test_in_venv_path_passes_parallel_to_run_server (below) is the runtime +# equivalent of the retired source-text guard for hardcoded +# `llama_parallel_slots = 4`. + + +# Re-exec arg-builder coverage. run() re-execs into the studio venv +# (execvp on POSIX, Popen on Windows). Without explicit forwarding the +# child reverts to typer defaults and silently drops the user's value. + + +class _ExecCaptured(SystemExit): + def __init__(self, argv): + super().__init__(0) + self.argv = list(argv) + + +def _install_reexec_capture(monkeypatch, *, platform): + studio_mod = _load_run_command() + captured = [] + + monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") + + fake_venv = Path("/fake/studio/venv/unsloth_studio") + fake_python = fake_venv / "bin" / "python" + fake_bin = fake_venv / "bin" / "unsloth" + monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_python) + + real_is_file = Path.is_file + monkeypatch.setattr( + Path, + "is_file", + lambda self: True if str(self) == str(fake_bin) else real_is_file(self), + ) + + # resolve_tool_policy is imported lazily inside run(); patch the source. + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + + monkeypatch.setattr(sys, "platform", platform) + + def fake_execvp(file, argv): + captured.append({"kind": "execvp", "argv": list(argv)}) + raise _ExecCaptured(argv) + + class _FakePopen: + def __init__(self, argv, *a, **kw): + captured.append({"kind": "popen", "argv": list(argv)}) + self._argv = argv + + def wait(self): + raise _ExecCaptured(self._argv) + + monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp) + monkeypatch.setattr(studio_mod.subprocess, "Popen", _FakePopen) + + return captured + + +def _invoke_run(monkeypatch, args, *, platform = "linux"): + import typer as _typer + + studio_mod = _load_run_command() + captured = _install_reexec_capture(monkeypatch, platform = platform) + app = _typer.Typer() + app.command( + context_settings = { + "allow_extra_args": True, + "ignore_unknown_options": True, + }, + )(studio_mod.run) + result = CliRunner().invoke(app, args, catch_exceptions = True) + return result, captured + + +def _value_after(argv, flag): + for i, tok in enumerate(argv): + if tok == flag and i + 1 < len(argv): + return argv[i + 1] + return None + + +_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"] + + +@pytest.mark.parametrize( + "flag,value", + [("--parallel", "8"), ("--n-parallel", "16"), ("-np", "32")], +) +def test_reexec_forwards_parallel_all_aliases(monkeypatch, flag, value): + """Every alias the user can type must reach the re-exec'd child.""" + result, captured = _invoke_run(monkeypatch, _BASE + [flag, value]) + assert ( + len(captured) == 1 + ), f"expected one launch via re-exec, got {captured}; output={result.output!r}" + argv = captured[0]["argv"] + assert ( + _value_after(argv, "--parallel") == value + ), f"{flag} {value} was dropped on re-exec; argv = {argv}" + + +@pytest.mark.parametrize("platform", ["linux", "darwin", "win32"]) +def test_reexec_argv_is_consistent_across_platforms(monkeypatch, platform): + """Linux/Darwin (execvp) and Windows (Popen) must build the same argv.""" + result, captured = _invoke_run( + monkeypatch, _BASE + ["--parallel", "12"], platform = platform + ) + assert len(captured) == 1 + expected_kind = "popen" if platform == "win32" else "execvp" + assert ( + captured[0]["kind"] == expected_kind + ), f"{platform}: expected launcher {expected_kind}, got {captured[0]['kind']}" + assert _value_after(captured[0]["argv"], "--parallel") == "12" + + +def test_reexec_np_is_first_class_alias(monkeypatch): + """`-np` must reach the child as --parallel . Pre-PR Click + clustered `-np 8` as `-p 8` (port=8) + stray `-n`; also pin that + --port is no longer collateral damage.""" + result, captured = _invoke_run(monkeypatch, _BASE + ["-np", "8"]) + assert len(captured) == 1 + argv = captured[0]["argv"] + assert ( + _value_after(argv, "--parallel") == "8" + ), f"-np 8 silently became 4 after re-exec; argv = {argv}" + # `-np 8` must not clobber --port (default 8888). + assert _value_after(argv, "--port") == "8888", argv + + +def test_reexec_mixed_parallel_with_passthrough(monkeypatch): + """--parallel + llama-server pass-through flags must all reach the child.""" + result, captured = _invoke_run( + monkeypatch, + _BASE + ["--parallel", "8", "--top-k", "20", "--temp", "0.7"], + ) + assert len(captured) == 1 + argv = captured[0]["argv"] + assert _value_after(argv, "--parallel") == "8", argv + assert _value_after(argv, "--top-k") == "20", argv + assert _value_after(argv, "--temp") == "0.7", argv + + +@pytest.mark.parametrize( + "user_flag,expected_in_child", + [ + ("--load-in-4bit", "--load-in-4bit"), + ("--no-load-in-4bit", "--no-load-in-4bit"), + (None, "--load-in-4bit"), # default True + ], +) +def test_reexec_forwards_load_in_4bit_in_both_directions( + monkeypatch, user_flag, expected_in_child +): + """Re-exec must emit the chosen polarity (or the typer default), + so a future default flip on one layer can't silently invert + behaviour for users who never typed the flag.""" + extras = [user_flag] if user_flag else [] + result, captured = _invoke_run(monkeypatch, _BASE + extras) + assert len(captured) == 1 + argv = captured[0]["argv"] + other_polarity = ( + "--no-load-in-4bit" + if expected_in_child == "--load-in-4bit" + else "--load-in-4bit" + ) + assert ( + expected_in_child in argv + ), f"expected {expected_in_child} in child argv; got {argv}" + assert ( + other_polarity not in argv + ), f"unexpected {other_polarity} in child argv; got {argv}" + + +# Runtime check: fake sys.prefix into the studio venv to bypass +# re-exec, then assert run_server receives --parallel as +# llama_parallel_slots. + + +class _RunServerCaptured(SystemExit): + def __init__(self, kwargs): + super().__init__(0) + self.kwargs = dict(kwargs) + + +def _types_module(name): + import types as _types + + return _types.ModuleType(name) + + +def test_studio_default_rejects_parallel_when_subcommand_invoked(): + """`unsloth studio --parallel 8 run ...` would silently drop the 8 + (typer doesn't forward parent options to subcommands). The + callback rejects with exit 2 and points at the subcommand flag.""" + studio_mod = _load_run_command() + import typer as _typer + + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + + runner = CliRunner() + result = runner.invoke(app, ["studio", "--parallel", "8", "run", "--model", "X"]) + assert result.exit_code == 2, ( + f"expected exit 2 when --parallel is on studio group with a " + f"subcommand invoked; got {result.exit_code}; output={result.output!r}" + ) + combined = (result.output or "") + (getattr(result, "stderr", "") or "") + assert "--parallel" in combined, combined + assert ( + "run --parallel 8" in combined + ), f"error message must show the corrected invocation; got: {combined}" + + +def test_studio_default_default_parallel_with_subcommand_does_not_error(): + """Omitting --parallel on the group must still let subcommands + run; the group's default 1 is benign.""" + studio_mod = _load_run_command() + import typer as _typer + + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + runner = CliRunner() + result = runner.invoke(app, ["studio", "--help"]) + assert result.exit_code == 0, result.output + + +def test_studio_default_exposes_parallel_option(): + """Plain `unsloth studio` exposes --parallel too so the API-only + path can raise concurrency without going through the denied + pass-through. Default stays at 1 (pre-PR); `run` keeps its 4.""" + studio_mod = _load_run_command() + import inspect + + sig = inspect.signature(studio_mod.studio_default) + assert "parallel" in sig.parameters, ( + "studio_default missing `parallel`; API-only path can't set " + "llama_parallel_slots" + ) + opt = sig.parameters["parallel"].default + decls = set(getattr(opt, "param_decls", []) or []) + assert "--parallel" in decls + assert "--n-parallel" in decls + assert ( + getattr(opt, "default", None) == 1 + ), "studio_default --parallel must default to 1 (pre-PR); `run` is 4" + assert getattr(opt, "min", None) == 1 + assert getattr(opt, "max", None) == 64 + + +@pytest.mark.parametrize("value", [1, 4, 8, 64]) +def test_in_venv_path_passes_parallel_to_run_server(monkeypatch, value): + """In-venv path must forward --parallel to + run_server(llama_parallel_slots=N), not the old hardcoded 4.""" + studio_mod = _load_run_command() + + fake_venv = Path("/fake/studio/venv/unsloth_studio") + monkeypatch.setattr(sys, "prefix", str(fake_venv)) + # Pin STUDIO_HOME so sys.prefix.startswith() picks the in-venv branch. + monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent) + + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + + captured: dict = {} + + def fake_run_server(**kwargs): + captured.update(kwargs) + raise _RunServerCaptured(kwargs) + + fake_backend_run = sys.modules.setdefault( + "studio.backend.run", _types_module("studio.backend.run") + ) + fake_backend_run.run_server = fake_run_server + fake_backend_run._resolve_external_ip = lambda: "127.0.0.1" + + import typer as _typer + + app = _typer.Typer() + app.command( + context_settings = { + "allow_extra_args": True, + "ignore_unknown_options": True, + }, + )(studio_mod.run) + CliRunner().invoke(app, _BASE + ["--parallel", str(value)], catch_exceptions = True) + + assert ( + captured.get("llama_parallel_slots") == value + ), f"run_server got llama_parallel_slots={captured.get('llama_parallel_slots')!r}, expected {value}" diff --git a/unsloth_cli/tests/test_studio_run_short_alias_clashes.py b/unsloth_cli/tests/test_studio_run_short_alias_clashes.py new file mode 100644 index 0000000000..8a3e94db47 --- /dev/null +++ b/unsloth_cli/tests/test_studio_run_short_alias_clashes.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for short-alias clashes with llama-server flags. + +`unsloth studio run` passes unknown flags through to llama-server. +Pre-cleanup it exposed 1-char shorts ``-m`` / ``-f`` plus ``-hfr``; +Click clustered llama-server tokens against them (``-fa`` -> ``-f a``, +``-mg 0`` -> ``-m g``, ``-fitt 1024`` -> ``-f itt``, ...), silently +breaking ~11 pass-through flags. + +The cleanup drops ``-m``, ``-f``, ``-hfr``. The 2-char ``-hf`` stays +(documented; multi-char shorts don't cluster). Long forms remain. +``studio_default`` keeps ``-f`` because it has no pass-through. + +See ``test_studio_run_parallel_flag.py`` for ``--parallel`` / +``-np`` coverage and re-exec forwarding. +""" + +from __future__ import annotations + +import inspect +import sys +from pathlib import Path + +import pytest +from typer.testing import CliRunner + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _studio_mod(): + from unsloth_cli.commands import studio as _s + + return _s + + +def _decls_for(param_name): + sig = inspect.signature(_studio_mod().run) + opt = sig.parameters[param_name].default + return set(getattr(opt, "param_decls", []) or []) + + +# Surface checks: removed shorts must not reappear. + + +def test_model_short_aliases_removed(): + """`-m` / `-hfr` removed from --model; `-hf` kept (multi-char, + doesn't cluster).""" + decls = _decls_for("model") + assert "-m" not in decls, "`-m` re-added; brings back `-mg`/`-md` clustering" + assert "-hfr" not in decls, "`-hfr` was re-added; remove it" + assert "--model" in decls + assert "--hf-repo" in decls + assert "-hf" in decls, "`-hf` is documented and must keep working" + + +def test_frontend_short_alias_removed_from_run(): + """`-f` must not be on `run` (eats `-fa`/`-fit`/`-fitt`/`-fitc`).""" + decls = _decls_for("frontend") + assert "-f" not in decls, "`-f` re-added on run(); brings back `-fa` clustering" + assert "--frontend" in decls + + +def test_studio_default_keeps_dash_f(): + """`studio_default` keeps `-f`: no pass-through tail to clash with.""" + sig = inspect.signature(_studio_mod().studio_default) + opt = sig.parameters["frontend"].default + decls = set(getattr(opt, "param_decls", []) or []) + assert "-f" in decls + + +# Behaviour checks: llama-server shorts must reach the child verbatim. + + +class _ExecCaptured(SystemExit): + def __init__(self, argv): + super().__init__(0) + self.argv = list(argv) + + +def _install_capture(monkeypatch): + studio_mod = _studio_mod() + captured = [] + monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") + fake_bin = Path("/fake/studio/venv/unsloth_studio/bin/unsloth") + monkeypatch.setattr( + studio_mod, "_studio_venv_python", lambda: fake_bin.parent / "python" + ) + real_is_file = Path.is_file + monkeypatch.setattr( + Path, + "is_file", + lambda self: True if str(self) == str(fake_bin) else real_is_file(self), + ) + from unsloth_cli import _tool_policy as _tp + + monkeypatch.setattr( + _tp, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + monkeypatch.setattr(sys, "platform", "linux") + + def fake_execvp(file, argv): + captured.append(list(argv)) + raise _ExecCaptured(argv) + + monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp) + return captured + + +def _invoke(monkeypatch, args): + import typer as _typer + + studio_mod = _studio_mod() + captured = _install_capture(monkeypatch) + app = _typer.Typer() + app.command( + context_settings = { + "allow_extra_args": True, + "ignore_unknown_options": True, + }, + )(studio_mod.run) + CliRunner().invoke(app, args, catch_exceptions = True) + return captured + + +# (short_flag, value, llama-server long name). All were silently +# mis-parsed pre-cleanup. +_PREVIOUSLY_BROKEN = [ + ("-fa", None, "--flash-attn"), + ("-fit", None, "--fit"), + ("-fitt", "1024", "--fit-target"), + ("-fitc", "4096", "--fit-ctx"), + ("-mg", "0", "--main-gpu"), + ("-md", "/path/draft.gguf", "--spec-draft-model"), + ("-hff", "Q4_K_M.gguf", "--hf-file"), + ("-cmoe", None, "--cpu-moe"), + ("-cram", "16384", "--cache-ram"), + ("-sm", "row", "--split-mode"), + ("-ncmoe", "8", "--n-cpu-moe"), +] + + +@pytest.mark.parametrize("flag,value,llama_long_name", _PREVIOUSLY_BROKEN) +def test_previously_broken_short_flag_now_passes_through( + monkeypatch, + flag, + value, + llama_long_name, +): + """Each of these was eaten by typer pre-cleanup; must pass through verbatim now.""" + extras = [flag] if value is None else [flag, value] + captured = _invoke(monkeypatch, ["--model", "X"] + extras) + assert len(captured) == 1, f"parent did not re-exec for {extras}" + argv = captured[0] + assert flag in argv, ( + f"llama-server short flag {flag!r} ({llama_long_name}) was eaten " + f"by typer; child argv = {argv}" + ) + if value is not None: + idx = argv.index(flag) + assert ( + idx + 1 < len(argv) and argv[idx + 1] == value + ), f"value for {flag!r} was lost or moved; argv = {argv}" + + +def test_dash_hf_documented_alias_still_works(monkeypatch): + """`-hf` is documented and must keep working (multi-char shorts + don't cluster in Click).""" + captured = _invoke( + monkeypatch, + ["-hf", "unsloth/gemma-4-26B-A4B-it-GGUF:UD-Q4_K_XL"], + ) + assert len(captured) == 1 + argv = captured[0] + # `_split_repo_variant` peels the `:variant` suffix before re-exec. + assert argv[argv.index("--model") + 1] == ("unsloth/gemma-4-26B-A4B-it-GGUF"), argv + assert argv[argv.index("--gguf-variant") + 1] == "UD-Q4_K_XL", argv + + +# Legacy `-m` / `-hfr` / `-f` were typer aliases pre-PR. The +# preprocessor promotes EXACT matches back to their typer params and +# leaves clustered tokens (`-mg`, `-fa`, ...) in the pass-through tail. + + +@pytest.mark.parametrize( + "legacy_args,expected_model", + [ + (["-m", "unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + (["-m=unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + (["-hfr", "unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + (["-hfr=unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + ], +) +def test_legacy_model_aliases_still_promote_to_model( + monkeypatch, + legacy_args, + expected_model, +): + """Pre-PR `-m X` / `-hfr X` set --model X; preprocessor preserves that.""" + captured = _invoke(monkeypatch, legacy_args) + assert len(captured) == 1, f"parent did not re-exec for {legacy_args}" + argv = captured[0] + assert argv[argv.index("--model") + 1] == expected_model, argv + # Promoted alias must not also leak into the pass-through tail. + for alias in ("-m", "-hfr"): + if alias in legacy_args: + assert alias not in argv, f"legacy {alias} leaked into child argv: {argv}" + + +def test_legacy_frontend_alias_still_promotes_to_frontend(monkeypatch): + """Pre-PR `-f dist` set --frontend dist; preprocessor preserves it.""" + captured = _invoke(monkeypatch, ["--model", "X", "-f", "/tmp/dist"]) + assert len(captured) == 1 + argv = captured[0] + # Compare via Path so Windows's str(Path("/tmp/dist")) = "\tmp\dist" + # doesn't trip the assertion on the same logical path. + assert Path(argv[argv.index("--frontend") + 1]) == Path("/tmp/dist"), argv + assert "-f" not in argv, f"-f leaked into child argv: {argv}" + + +def test_legacy_model_alias_conflicts_with_long_form(monkeypatch): + """`--model X` plus `-m Y` is ambiguous; must error pre-re-exec.""" + captured = _invoke(monkeypatch, ["--model", "X", "-m", "Y"]) + assert ( + len(captured) == 0 + ), f"expected error before re-exec, got launch with argv = {captured}" + + +def test_clustered_tokens_are_not_promoted(monkeypatch): + """`-mg` / `-fa` / `-fitt` are llama-server flags and must survive + in the tail even though they start with `-m` / `-f`.""" + captured = _invoke( + monkeypatch, + ["--model", "X", "-mg", "0", "-fa", "-fitt", "1024"], + ) + assert len(captured) == 1 + argv = captured[0] + assert argv[argv.index("--model") + 1] == "X", argv + for flag in ("-mg", "-fa", "-fitt"): + assert flag in argv, f"{flag!r} was promoted instead of passed through: {argv}" + + +def test_legacy_m_with_repo_variant_syntax(monkeypatch): + """`-m repo:variant` must round-trip through preprocessor + + _split_repo_variant into --model + --gguf-variant.""" + captured = _invoke( + monkeypatch, + ["-m", "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL"], + ) + assert len(captured) == 1 + argv = captured[0] + assert argv[argv.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF", argv + assert argv[argv.index("--gguf-variant") + 1] == "UD-Q4_K_XL", argv + + +def test_missing_model_after_preprocessor_errors(monkeypatch): + """Neither --model nor a legacy alias → clean exit(2) before re-exec.""" + captured = _invoke(monkeypatch, ["--parallel", "8"]) + assert ( + len(captured) == 0 + ), f"expected exit before re-exec, got launch with argv = {captured}" + + +def test_legacy_m_inline_value_form(monkeypatch): + """`-m=foo` is promoted like `-m foo`.""" + captured = _invoke(monkeypatch, ["-m=unsloth/Qwen3-1.7B-GGUF"]) + assert len(captured) == 1 + argv = captured[0] + assert argv[argv.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF", argv + + +# Unit tests for _consume_legacy_short_aliases. + + +def test_consume_helper_exact_match_space_form(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-m", "FOO", "--top-k", "20"], + ("-m",), + None, + "--model", + ) + assert value == "FOO" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_exact_match_inline_form(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-m=FOO", "--top-k", "20"], + ("-m",), + None, + "--model", + ) + assert value == "FOO" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_leaves_clusters_alone(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-mg", "0", "-md", "/x"], + ("-m",), + None, + "--model", + ) + assert value is None + assert remaining == ["-mg", "0", "-md", "/x"] + + +def test_consume_helper_value_already_set_raises(): + helper = _studio_mod()._consume_legacy_short_aliases + import typer as _typer + + with pytest.raises(_typer.BadParameter): + helper(["-m", "Y"], ("-m",), "X", "--model") + + +def test_consume_helper_missing_value_raises(): + helper = _studio_mod()._consume_legacy_short_aliases + import typer as _typer + + with pytest.raises(_typer.BadParameter): + helper(["-m"], ("-m",), None, "--model") + + +def test_consume_helper_multiple_aliases_in_group(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-hfr", "FOO", "--top-k", "20"], + ("-m", "-hfr"), + None, + "--model", + ) + assert value == "FOO" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_preserves_value_when_no_match(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["--top-k", "20"], + ("-m",), + "PRESET", + "--model", + ) + assert value == "PRESET" + assert remaining == ["--top-k", "20"] + + +# `-p` is typer short for --port, so Click clusters `-np8` as `-n -p 8` +# (port=8, parallel dropped). The rewrite splits to `-np 8` pre-parse. + + +def test_expand_np_rewrites_attached_form(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "studio", "run", "--model", "X", "-np8"], + ) + _studio_mod()._expand_attached_np_short() + assert sys.argv == [ + "unsloth", + "studio", + "run", + "--model", + "X", + "-np", + "8", + ] + + +@pytest.mark.parametrize("value", ["1", "8", "64", "999"]) +def test_expand_np_rewrites_all_digit_values(monkeypatch, value): + monkeypatch.setattr(sys, "argv", ["unsloth", "studio", "run", f"-np{value}"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "studio", "run", "-np", value] + + +def test_expand_np_leaves_space_form_alone(monkeypatch): + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-np", "8"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", "8"] + + +def test_expand_np_leaves_equals_form_alone(monkeypatch): + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-np=8"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np=8"] + + +def test_expand_np_leaves_non_digit_suffix_alone(monkeypatch): + # `-npfoo` isn't a numeric attached value; let typer reject it. + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-npfoo"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-npfoo"] + + +def test_expand_np_leaves_bare_np_alone(monkeypatch): + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-np"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np"] + + +def test_expand_np_handles_multiple_occurrences(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "run", "-np8", "-np16"], + ) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", "8", "-np", "16"] + + +@pytest.mark.parametrize("attached,expected", [("-np-1", "-1"), ("-np+1", "+1")]) +def test_expand_np_handles_signed_attached_forms(monkeypatch, attached, expected): + """Signed `-np-1` / `-np+1` must split too, else Click sets port=-1.""" + monkeypatch.setattr(sys, "argv", ["unsloth", "run", attached]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", expected] + + +@pytest.mark.parametrize( + "attached,expected_suffix", + [("-np8x", "8x"), ("-np-1foo", "-1foo"), ("-np9bar", "9bar")], +) +def test_expand_np_rewrites_numeric_prefix_even_with_junk( + monkeypatch, attached, expected_suffix +): + """`-np8x` would surface as a baffling --port error; rewriting to + `-np 8x` makes typer report against `-np` where it was typed.""" + monkeypatch.setattr(sys, "argv", ["unsloth", "run", attached]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", expected_suffix] + + +def test_consume_helper_rejects_empty_inline_value(): + """`-m=` must error, not silently become --model ''.""" + import typer as _typer + + helper = _studio_mod()._consume_legacy_short_aliases + with pytest.raises(_typer.BadParameter, match = "non-empty"): + helper(["-m="], ("-m",), None, "--model") + + +# Gate isolation: importing unsloth_cli from a third-party script must +# leave its sys.argv intact. Pins the narrow basename set. + + +@pytest.mark.parametrize( + "third_party_argv0", + [ + "/home/user/myproj/cli.py", + "cli.py", + "/usr/bin/some-tool", + "pytest", + "/opt/wrapper/launch.py", + "unsloth-cli", + "unsloth-cli.py", + ], +) +def test_third_party_importers_do_not_trigger_np_rewrite( + monkeypatch, third_party_argv0 +): + """Only the `unsloth` / `unsloth.exe` console-script may run the + canonicaliser; third-party scripts must keep their argv intact.""" + import os as _os + import importlib + + starting_argv = [third_party_argv0, "subcmd", "-np8", "--input", "foo"] + monkeypatch.setattr(sys, "argv", list(starting_argv)) + # Force a fresh import so the import-time gate actually runs. + monkeypatch.delitem(sys.modules, "unsloth_cli", raising = False) + importlib.import_module("unsloth_cli") + assert sys.argv == starting_argv, ( + f"third-party argv[0]={third_party_argv0!r} triggered the " + f"unsloth -np canonicaliser; sys.argv was mutated to {sys.argv}" + ) + _ = _os # silence unused-import linters when monkeypatch lazy-binds + + +def test_attached_np8_no_longer_silently_sets_port(monkeypatch): + """After the gate runs, `-np8` produces --parallel=8 (not --port=8).""" + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "studio", "run", "--model", "X", "-np8"], + ) + _studio_mod()._expand_attached_np_short() + captured = _invoke(monkeypatch, sys.argv[2:]) # drop "unsloth studio" + assert len(captured) == 1, "parent did not re-exec" + argv = captured[0] + assert argv[argv.index("--parallel") + 1] == "8", argv + assert argv[argv.index("--port") + 1] == "8888", argv + + +def test_expand_np_stops_at_double_dash(monkeypatch): + """Tokens after `--` are positional; `-np8` stays raw.""" + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "run", "--model", "X", "--", "-np8"], + ) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "--model", "X", "--", "-np8"] + + +def test_consume_helper_stops_at_double_dash(): + """Alias promotion must not reach past `--`.""" + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["--top-k", "20", "--", "-m", "FOO"], + ("-m",), + None, + "--model", + ) + assert value is None + assert remaining == ["--top-k", "20", "--", "-m", "FOO"] + + +def test_consume_helper_rejects_long_flag_as_value(): + """`-m --flash-attn` errors; `--xxx` is unambiguously a flag.""" + import typer as _typer + + helper = _studio_mod()._consume_legacy_short_aliases + with pytest.raises(_typer.BadParameter, match = "--flash-attn"): + helper(["-m", "--flash-attn"], ("-m",), None, "--model") + + +def test_consume_helper_allows_bare_dash_as_value(): + """Lone `-` is a stdin/path sentinel, not a flag.""" + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper(["-m", "-", "--top-k", "20"], ("-m",), None, "--model") + assert value == "-" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_allows_short_dash_value(): + """`-foo` may be a path or a leading-dash model name; only `--long` + tokens are rejected as values.""" + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper(["-m", "-foo", "--top-k", "20"], ("-m",), None, "--model") + assert value == "-foo" + assert remaining == ["--top-k", "20"] From 02ea6c9233c8d86670fb1d3b7fe97b9edb58ee4b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 00:30:30 -0700 Subject: [PATCH 29/43] tests: unblock three stale assertions broken on main (MLX CI + Backend CI) (#5803) * tests: unblock three stale assertions broken on main MLX CI on Mac M1 + Backend CI (both Repo tests CPU and Python 3.10/11/12/13) have been red on every push to main for days. None of the underlying code is wrong; three test files have stale anchors / assertions left behind by PR #5537 (max_steps bump) and PR #5775 (composer + provision-desktop-auth). 1. tests/studio/run_real_mlx_smoke.py:393 PR #5537 bumped max_steps from 7 to 30 for seed-robust convergence but left `assert len(losses_per_step) == 7`. With logging_steps=1 the callback fires once per step; 30 entries, not 7. Track config.max_steps so the gate auto-follows future bumps. 2. tests/studio/test_composer_rtl_bidi_attribute.py:29 PR #5775 changed the composer aria-label from the literal `aria-label="Message input"` to a JSX ternary `aria-label={overlay ? "Image edit instructions" : "Message input"}`. Anchor on the inner string literal `"Message input"` instead. 3. studio/backend/tests/test_desktop_auth.py:487 The guarded_import in test_provision_desktop_auth_writes_secret_and_creates_db_without_backend_deps blocks any import whose name == "utils", including the relative `from .utils import echo` inside typer._click.decorators (typer 0.25+). Gate the block on level == 0 so only absolute imports of `utils` / `auth` / `fastapi` / `structlog` are rejected; relative imports inside third-party packages pass through. All three tests pass locally; the MLX one is a mechanical 7->config.max_steps swap and will be exercised by MLX CI on this PR. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/tests/test_desktop_auth.py | 9 ++++++--- tests/studio/run_real_mlx_smoke.py | 6 +++++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index ab1a03eeda..bc8788fd90 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -484,11 +484,14 @@ from typer.testing import CliRunner studio_home = Path(sys.argv[1]) real_import = builtins.__import__ -def guarded_import(name, *args, **kwargs): +def guarded_import(name, globals = None, locals = None, fromlist = (), level = 0): + # Only gate absolute imports; relative `from .utils import x` inside + # third-party packages (e.g. typer._click.decorators) hits level > 0 + # with name="utils" and must pass through. blocked = ("auth", "fastapi", "structlog", "utils") - if name in blocked or name.startswith(("auth.", "utils.")): + if level == 0 and (name in blocked or name.startswith(("auth.", "utils."))): raise ModuleNotFoundError(name) - return real_import(name, *args, **kwargs) + return real_import(name, globals, locals, fromlist, level) builtins.__import__ = guarded_import from unsloth_cli.commands import studio as studio_cli diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 27f682ee4e..dc3001fd99 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -390,7 +390,11 @@ def cmd_train(args) -> int: ) if k in train_result } - assert len(losses_per_step) == 7, f"expected 7 logged steps, got {losses_per_step}" + # logging_steps=1 + max_steps=N -> N callbacks; track config so the + # gate auto-follows if max_steps is bumped again. + assert ( + len(losses_per_step) == config.max_steps + ), f"expected {config.max_steps} logged steps, got {losses_per_step}" for i, l in enumerate(losses_per_step): # Allow exact 0.0: fp16 per-step loss underflows to 0.0 after # the LoRA reaches loss=0 around step ~10 with this fixture + From 556f396b3cd63da525e2cdac1a5f1c606f4494ce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 01:35:13 -0700 Subject: [PATCH 30/43] ci: install unsloth_zoo from git main in notebooks-ci + studio-backend-ci (#5802) * ci: install unsloth_zoo from git main in notebooks-ci + studio-backend-ci These were the only two workflows that still pulled unsloth_zoo from PyPI; every other CI (Core, MLX, version-compat, install.sh-driven Studio smokes) installs zoo from git main. Drift between PyPI and main hides fixes-on-zoo-main and lets PR-time validation pass on a stale zoo, then break for users on next release. Both edits match the retry-with-backoff shape mlx-ci.yml already uses. * ci: drop --no-deps from studio-backend-ci unsloth_zoo install The prior PyPI line was `pip install 'unsloth_zoo>=2026.5.1'` (no --no-deps), which pulled in triton and the rest of zoo's runtime deps. I dropped that transitive resolve in the first commit, which broke collection of 5 tests in Repo tests (CPU) with ModuleNotFoundError: No module named 'triton'. Match the prior dep-resolve shape, keeping the source-from-git change. notebooks-ci keeps --no-deps because its original line also had it. --- .github/workflows/notebooks-ci.yml | 10 +++++++++- .github/workflows/studio-backend-ci.yml | 16 +++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 673b2f3cc5..2edcae8ab2 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -285,7 +285,15 @@ jobs: # The PR-time CI must validate the code in this PR; PyPI unsloth # may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py # (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream. - pip install --no-deps unsloth_zoo + # unsloth_zoo from git main mirrors every other CI (Core / MLX / + # install.sh) so PR-time validation sees the same zoo HEAD. + for attempt in 1 2 3; do + if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install --no-deps -e ./unsloth - name: Convert notebooks for AST scan diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 63eb70f7f1..ee5bbe8633 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -144,9 +144,19 @@ jobs: # versions ship a CPU build that imports cleanly on Linux. pip install 'bitsandbytes>=0.45' # unsloth.device_type imports unsloth_zoo.utils.Version at module - # scope, so the conftest preload needs unsloth_zoo even though - # it is an optional dep of unsloth. - pip install 'unsloth_zoo>=2026.5.1' + # scope, so the conftest preload needs unsloth_zoo. Pull from + # git main so this job sees the same zoo HEAD as Core / MLX / + # install.sh do (otherwise a fix on zoo main hides until release). + # No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'` + # behaviour so triton etc. still come in for the Repo tests CPU + # collection imports. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install -e . --no-deps - name: Repo tests (CPU, auto-discovered) From 8b2b99be036dc114700867392f1b6a149ee3fdaf Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 27 May 2026 16:50:44 +0530 Subject: [PATCH 31/43] tool mask support (#5682) * tool mask support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle tool masks with older zoo builds * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep tool mask implementation in zoo --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_tool_mask_zoo_compat.py | 103 +++++++++++++++++++++++++++++ unsloth/models/rl.py | 19 ++++++ unsloth/models/rl_replacements.py | 45 ++++++++++++- 3 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 tests/test_tool_mask_zoo_compat.py diff --git a/tests/test_tool_mask_zoo_compat.py b/tests/test_tool_mask_zoo_compat.py new file mode 100644 index 0000000000..b6b68a7561 --- /dev/null +++ b/tests/test_tool_mask_zoo_compat.py @@ -0,0 +1,103 @@ +"""Compatibility checks for env/tool mask support with older unsloth_zoo.""" + +from __future__ import annotations + +import ast +import os +import textwrap + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)) +RL_SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl.py") +RL_REPLACEMENTS_SOURCE_PATH = os.path.join( + REPO_ROOT, "unsloth", "models", "rl_replacements.py" +) + + +def _read(path: str) -> str: + with open(path, "r") as fh: + return fh.read() + + +def _load_local_align_completion_tool_mask(): + src = _read(RL_SOURCE_PATH) + tree = ast.parse(src) + for node in tree.body: + if isinstance(node, ast.If): + for item in node.body: + if ( + isinstance(item, ast.FunctionDef) + and item.name == "align_completion_tool_mask" + ): + function_src = ast.get_source_segment(src, item) + break + else: + continue + break + else: + raise AssertionError("local align_completion_tool_mask fallback is missing") + + calls = [] + + def align_logprobs_with_mask(logprob_tensor, completion_mask, pad_value = None): + calls.append((logprob_tensor, completion_mask, pad_value)) + return torch.tensor( + [[1, 0, 1], [0, 1, 1]], + device = completion_mask.device, + dtype = logprob_tensor.dtype, + ) + + namespace = { + "torch": torch, + "align_logprobs_with_mask": align_logprobs_with_mask, + } + exec(textwrap.dedent(function_src), namespace) + return namespace["align_completion_tool_mask"], calls + + +def test_rl_uses_optional_zoo_tool_mask_helper(): + src = _read(RL_SOURCE_PATH) + assert 'RL_REPLACEMENTS.get("align_completion_tool_mask")' in src + assert 'RL_REPLACEMENTS["align_completion_tool_mask"]' not in src + + +def test_local_tool_mask_fallback_is_only_old_zoo_compat_shim(): + align_completion_tool_mask, calls = _load_local_align_completion_tool_mask() + completion_mask = torch.tensor( + [[1, 1, 0], [1, 1, 1]], + dtype = torch.float32, + ) + + assert align_completion_tool_mask(None, completion_mask) is completion_mask + assert calls == [] + + same_shape_tool_mask = torch.tensor([[1, 0, 1], [0, 1, 1]], dtype = torch.bool) + with pytest.raises(RuntimeError, match = "Please upgrade unsloth_zoo"): + align_completion_tool_mask(same_shape_tool_mask, completion_mask) + + +def test_grpo_accumulated_loss_omits_none_tool_mask_for_old_zoo(): + src = _read(RL_REPLACEMENTS_SOURCE_PATH) + assert "_grpo_accumulated_loss_kwargs = {}" in src + assert ( + 'if tool_mask is not None:\n _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask' + in src + ) + assert src.count("**_grpo_accumulated_loss_kwargs") == 2 + + accelerated_loss_start = src.find('if hasattr(self.args, "loss_type"):') + assert accelerated_loss_start != -1 + accelerated_loss_body = src[ + accelerated_loss_start : src.find( + 'if "train" in self._metrics:', accelerated_loss_start + ) + ] + assert "tool_mask = tool_mask" not in accelerated_loss_body + + +def test_rollout_output_patch_requires_real_tool_mask_symbol(): + src = _read(RL_REPLACEMENTS_SOURCE_PATH) + assert 're.search(r"\\btool_mask\\b", function)' in src + assert 'output["tool_mask"]' in src diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index c82c8364b3..9bf0d3e968 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -365,6 +365,22 @@ calculate_pad_tokens_in_prompt = RL_REPLACEMENTS["calculate_pad_tokens_in_prompt create_completion_attention_mask = RL_REPLACEMENTS["create_completion_attention_mask"] left_pack_padding = RL_REPLACEMENTS["left_pack_padding"] align_logprobs_with_mask = RL_REPLACEMENTS["align_logprobs_with_mask"] +align_completion_tool_mask = RL_REPLACEMENTS.get("align_completion_tool_mask") +if align_completion_tool_mask is None: + + def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, + ) -> torch.Tensor: + if tool_mask is None: + return completion_mask + raise RuntimeError( + "env_mask/tool_mask GRPO requires an unsloth_zoo build whose " + "grpo_accumulated_loss handles tool_mask. Please upgrade " + "unsloth_zoo." + ) + + autotune_batch_and_chunks = RL_REPLACEMENTS["grpo_autotune_batch_and_chunks"] sanitize_logprob = RL_REPLACEMENTS["sanitize_logprob"] @@ -452,6 +468,7 @@ torch_compile_options = {{ {create_completion_attention_mask_code} {left_pack_padding_code} {align_logprobs_with_mask_code} +{align_completion_tool_mask_code} {autotune_batch_and_chunks_code} {sanitize_logprob_code} @@ -1577,6 +1594,7 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): ) left_pack_padding_code = inspect.getsource(left_pack_padding) align_logprobs_with_mask_code = inspect.getsource(align_logprobs_with_mask) + align_completion_tool_mask_code = inspect.getsource(align_completion_tool_mask) autotune_batch_and_chunks_code = inspect.getsource(autotune_batch_and_chunks) sanitize_logprob_code = inspect.getsource(sanitize_logprob) # Get final source code @@ -1607,6 +1625,7 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): autotune_batch_and_chunks_code = autotune_batch_and_chunks_code, left_pack_padding_code = left_pack_padding_code, align_logprobs_with_mask_code = align_logprobs_with_mask_code, + align_completion_tool_mask_code = align_completion_tool_mask_code, sanitize_logprob_code = sanitize_logprob_code, ) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 0f9a324d5b..9ddb5e9453 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -946,6 +946,14 @@ def grpo_trainer__generate_and_score_completions(function_name, function): ) function = function.replace(_save_search, _save_replace) + if re.search(r"\btool_mask\b", function) and 'output["tool_mask"]' not in function: + function = function.replace( + " return output", + " if tool_mask is not None:\n" + ' output["tool_mask"] = tool_mask\n' + " return output", + ) + return function @@ -1523,6 +1531,7 @@ def grpo_trainer_compute_loss(function_name, function): mm_token_type_ids = inputs.get("mm_token_type_ids", None) num_items_in_batch = inputs.get("num_items_in_batch", None) sampling_per_token_logps = inputs.get("sampling_per_token_logps", None) + tool_mask = inputs.get("tool_mask", None) current_gradient_accumulation_steps = self.current_gradient_accumulation_steps num_processes = self.accelerator.num_processes @@ -1598,6 +1607,16 @@ def grpo_trainer_compute_loss(function_name, function): max_left_pad = inputs.get("max_left_pad", 0) if per_token_logps is not None: + loss_mask = completion_mask + if tool_mask is not None: + if tool_mask.shape != completion_mask.shape: + raise ValueError( + "tool_mask/env_mask must have the same shape as completion_mask" + ) + loss_mask = completion_mask * tool_mask.to( + device = completion_mask.device, + dtype = completion_mask.dtype, + ) ( loss, completion_length, @@ -1612,7 +1631,7 @@ def grpo_trainer_compute_loss(function_name, function): old_logps, sampling_per_token_logps, input_ids, - completion_mask, + loss_mask, self.beta, advantages, pixel_values = pixel_values, @@ -1662,6 +1681,28 @@ def grpo_trainer_compute_loss(function_name, function): "unsloth_zoo (see https://github.com/unslothai/unsloth-zoo/pull/613)." ) self._unsloth_grpo_zoo_checked = True + if tool_mask is not None and not getattr( + self, "_unsloth_grpo_tool_mask_zoo_checked", False + ): + _supports_tool_mask = ( + "tool_mask" in inspect.signature(grpo_accumulated_loss).parameters + ) + if not _supports_tool_mask: + try: + _zoo_src = inspect.getsource(grpo_accumulated_loss) + except (TypeError, OSError): + _zoo_src = "" + _supports_tool_mask = "tool_mask" in _zoo_src + if not _supports_tool_mask: + raise RuntimeError( + "env_mask/tool_mask GRPO requires an unsloth_zoo build whose " + "grpo_accumulated_loss handles tool_mask. Please upgrade " + "unsloth_zoo." + ) + self._unsloth_grpo_tool_mask_zoo_checked = True + _grpo_accumulated_loss_kwargs = {} + if tool_mask is not None: + _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask if hasattr(self.args, "loss_type"): ( loss, @@ -1703,6 +1744,7 @@ def grpo_trainer_compute_loss(function_name, function): sampling_per_token_logps = sampling_per_token_logps, token_type_ids = token_type_ids, mm_token_type_ids = mm_token_type_ids, + **_grpo_accumulated_loss_kwargs, ) else: # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 @@ -1728,6 +1770,7 @@ def grpo_trainer_compute_loss(function_name, function): attention_mask = attention_mask, token_type_ids = token_type_ids, mm_token_type_ids = mm_token_type_ids, + **_grpo_accumulated_loss_kwargs, ) ) if "train" in self._metrics: From 4891118b5e31194e79c361de5b6b49cefe9b4112 Mon Sep 17 00:00:00 2001 From: wuwuwu <1498968550@qq.com> Date: Wed, 27 May 2026 19:52:16 +0800 Subject: [PATCH 32/43] Studio: add frontend i18n support (#5765) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: add frontend i18n support * Studio i18n: guard storage events, restore plurals, fill zh-CN, add parity check - locale-store.ts: wrap window.localStorage access in handleStorageEvent with try/catch. readStoredLocale and writeStoredLocale already guard the same API; the storage-event path can throw the same way in privacy/ restricted contexts and was the only unguarded localStorage call. Refactor the storageArea + key match into isLocaleStorageEvent for clarity. - chat-tab.tsx + en.ts/zh-CN.ts: restore singular handling for chat-clear copy that the i18n migration dropped. Pre-PR code rendered "1 chat" but the new template strings always said "chats", so a user with exactly one chat saw "Cleared 1 chats", "Clear 1 chats?", and "1 chats cleared; 1 chats remain". Add clearOneChat*, clearedOneChat, oneChatClearedRemain*, chatsClearedRemainOne, and storageClearFailedOne keys and pick them in chat-tab.tsx when count === 1. - zh-CN.ts: fill ~50 previously English-fallback keys across studio.configure, studio.model VRAM helpers, studio.dataset (source, browsing, tooltips, preview/split/subset), studio.params tooltips and learningRateDescription, studio.training (audio/vision incompatible), studio.trainingStart.terminalStart, studio.tour.guidedTour, settings.chat.clear*, settings.connections, settings.apiKeys.newBadge. shell.{beta,brand,product} kept as brand strings. - src/i18n/check-parity.ts + npm i18n:check: small script that verifies every locale overlay against the English baseline. Catches placeholder mismatches, shape mismatches, and unintended extra keys; runs via node --experimental- strip-types with no new devDependencies. Verified locally: npm run typecheck, lint, build, biome:check, i18n:check all pass. 24 vitest unit tests cover locale resolution, persistence failures, storage-event sync (including window.localStorage throwing), interpolation, and fallback. 33 Playwright e2e tests pass across Chromium, Firefox, and WebKit covering default load, switch + reload persistence, unsupported/garbage locale fallback, storage-event cross-tab sync, and storage clear. * Studio i18n: use translated API-key error copy instead of raw err.message The API helpers in src/features/settings/api/api-keys.ts throw generic English Error objects ("Failed to load API access", "Failed to create access token", "Failed to revoke access token"). ApiKeysTab and CreateKeyForm caught those and preferred err.message over the translated "settings.apiKeys.loadError" / .createError / .revokeError keys, so in zh-CN mode failed load/create/revoke requests still surfaced the English strings instead of the translated copy. Switched the four call-sites to always render the translated message and left the helper throws unchanged (they are still useful for diagnostics but should not be treated as user-facing localized copy). * Studio i18n: polish two zh-CN embedding LR tooltips Translation-pass review surfaced two awkward phrasings I introduced earlier: "常用区间是主学习率的 2 至 10 倍小" -> "常用区间是比主学习率小 2 至 10 倍" Both versions are grammatical, but the new "比 X 小 N 倍" phrasing is the standard idiomatic comparative for "N times smaller than X" in technical Chinese writing. The earlier "X 的 N 倍小" reads as a non-native construction. Applies to: studio.params.embeddingLearningRateTooltip studio.params.embeddingLearningRateDescription --------- Co-authored-by: Daniel Han --- studio/frontend/package.json | 1 + studio/frontend/src/app/router.tsx | 9 +- studio/frontend/src/app/routes/__root.tsx | 39 +- studio/frontend/src/app/routes/studio.tsx | 2 +- .../frontend/src/components/app-sidebar.tsx | 162 ++-- .../profile-personalization-panel.tsx | 25 +- .../settings/components/api-key-row.tsx | 58 +- .../settings/components/create-key-form.tsx | 19 +- .../settings/components/key-reveal-card.tsx | 14 +- .../settings/components/language-select.tsx | 46 ++ .../settings/components/theme-segmented.tsx | 16 +- .../components/update-studio-instructions.tsx | 94 +-- .../settings/components/usage-examples.tsx | 24 +- .../src/features/settings/settings-dialog.tsx | 46 +- .../src/features/settings/tabs/about-tab.tsx | 34 +- .../features/settings/tabs/api-keys-tab.tsx | 72 +- .../features/settings/tabs/appearance-tab.tsx | 30 +- .../src/features/settings/tabs/chat-tab.tsx | 78 +- .../features/settings/tabs/general-tab.tsx | 56 +- .../features/settings/tabs/profile-tab.tsx | 9 +- .../studio/historical-training-view.tsx | 29 +- .../src/features/studio/history-card-grid.tsx | 93 ++- .../sections/charts/chart-settings-sheet.tsx | 60 +- .../sections/charts/eval-loss-chart-card.tsx | 26 +- .../sections/charts/grad-norm-chart-card.tsx | 19 +- .../charts/learning-rate-chart-card.tsx | 19 +- .../charts/training-loss-chart-card.tsx | 28 +- .../studio/sections/dataset-section.tsx | 139 ++-- .../studio/sections/model-section.tsx | 81 +- .../studio/sections/params-section.tsx | 212 ++--- .../studio/sections/progress-section-lib.ts | 13 - .../studio/sections/progress-section.tsx | 125 +-- .../studio/sections/training-section.tsx | 53 +- .../src/features/studio/studio-page.tsx | 23 +- .../studio/training-start-overlay.tsx | 45 +- .../frontend/src/features/training/index.ts | 5 +- studio/frontend/src/i18n/AGENTS.md | 11 + studio/frontend/src/i18n/check-parity.ts | 111 +++ studio/frontend/src/i18n/index.ts | 44 ++ studio/frontend/src/i18n/locale-store.ts | 130 ++++ studio/frontend/src/i18n/locales/en.ts | 731 ++++++++++++++++++ studio/frontend/src/i18n/locales/zh-CN.ts | 712 +++++++++++++++++ studio/frontend/src/i18n/messages.ts | 76 ++ studio/frontend/src/i18n/types.ts | 30 + studio/frontend/src/main.tsx | 3 + 45 files changed, 2958 insertions(+), 694 deletions(-) create mode 100644 studio/frontend/src/features/settings/components/language-select.tsx create mode 100644 studio/frontend/src/i18n/AGENTS.md create mode 100644 studio/frontend/src/i18n/check-parity.ts create mode 100644 studio/frontend/src/i18n/index.ts create mode 100644 studio/frontend/src/i18n/locale-store.ts create mode 100644 studio/frontend/src/i18n/locales/en.ts create mode 100644 studio/frontend/src/i18n/locales/zh-CN.ts create mode 100644 studio/frontend/src/i18n/messages.ts create mode 100644 studio/frontend/src/i18n/types.ts diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 83b1fd96f9..b43e174889 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -12,6 +12,7 @@ "lint": "eslint .", "preview": "vite preview", "typecheck": "tsc -b --pretty false", + "i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts", "biome:check": "biome check", "biome:fix": "biome check --write" }, diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index c7bc0440bd..f0a417638d 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -3,6 +3,7 @@ import { Link, createRouter, useRouterState } from "@tanstack/react-router"; import { Button } from "@/components/ui/button"; +import { useT } from "@/i18n"; import { Route as rootRoute } from "./routes/__root"; import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; @@ -31,7 +32,9 @@ const routeTree = rootRoute.addChildren([ ]); function DefaultNotFound() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); + return (

- Page not found + {t("shell.notFound.title")}

- {pathname} does not exist. + {t("shell.notFound.description", { path: pathname })}

); diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 47bff815e6..57ed233d51 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,8 +6,9 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; -import { useTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard"; +import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; +import { useT, type TranslationKey } from "@/i18n"; import { Outlet, createRootRoute, @@ -16,24 +17,25 @@ import { useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; -import { Suspense, useEffect, useLayoutEffect, type ReactNode } from "react"; +import { Suspense, useEffect, useLayoutEffect } from "react"; import { AppProvider } from "../provider"; -// Type `staticData.title` on every route so the matched-title selector -// below stays type-safe without an inline cast. declare module "@tanstack/react-router" { interface StaticDataRouteOption { title?: string; + titleKey?: TranslationKey; } } -// Fallback while a lazy route bundle (Train/Recipes/Export) loads. -// /chat is synchronous and never hits this. -const RouteFallback: ReactNode = ( -
- Loading... -
-); +function RouteFallback() { + const t = useT(); + + return ( +
+ {t("common.loading")} +
+ ); +} const CHAT_ONLY_ALLOWED = new Set([ "/", @@ -68,6 +70,7 @@ const HIDDEN_NAVBAR_ROUTES = ["/onboarding", "/login", "/change-password"]; const DEFAULT_DOCUMENT_TITLE = "Unsloth Studio"; function RootLayout() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname); const isChatRoute = pathname.startsWith("/chat"); @@ -75,24 +78,20 @@ function RootLayout() { useTrainingUnloadGuard(); - // Walk matches deepest-first; each route declares its own title. const matchedTitle = useMatches({ select: (matches) => { for (let i = matches.length - 1; i >= 0; i--) { - const title = matches[i].staticData.title; + const { title, titleKey } = matches[i].staticData; + if (titleKey) return t(titleKey); if (title) return title; } return null; }, }); - // `/settings` redirects in `beforeLoad`, so its route never stays - // matched; surface the modal's title via the store instead. const settingsDialogOpen = useSettingsDialogStore((s) => s.open); - const documentTitle = settingsDialogOpen ? "Settings" : matchedTitle; + const documentTitle = settingsDialogOpen ? t("settings.title") : matchedTitle; - // useLayoutEffect updates the tab title before paint, avoiding a - // one-frame flash of the previous route's title on navigation. useLayoutEffect(() => { document.title = documentTitle ? `${documentTitle} - ${DEFAULT_DOCUMENT_TITLE}` @@ -116,7 +115,7 @@ function RootLayout() { {hideNavbar ? (
- + }>
@@ -142,7 +141,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"}`} > - + }> diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx index 75f1a1b937..ae7f445e94 100644 --- a/studio/frontend/src/app/routes/studio.tsx +++ b/studio/frontend/src/app/routes/studio.tsx @@ -15,7 +15,7 @@ const StudioPage = lazy(() => export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/studio", - staticData: { title: "Train" }, + staticData: { titleKey: "studio.routeTitle" }, beforeLoad: () => requireAuth(), component: StudioPage, }); diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index aac5f8f8a8..849e017ea8 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -90,9 +90,33 @@ import { useTrainingRuntimeStore, } from "@/features/training"; import type { TrainingRunSummary } from "@/features/training"; -import { useEffect, useRef, useState } from "react"; +import { useEffect, useRef, useState, type ReactNode } from "react"; import { toast } from "@/lib/toast"; import { ShutdownDialog } from "@/components/shutdown-dialog"; +import { translate, useT, type TranslationKey } from "@/i18n"; + +const EMPHASIS_MARKER = "__UNSLOTH_I18N_EMPHASIS_MARKER__"; + +type AppT = ReturnType; + +function renderEmphasizedTranslation( + t: AppT, + key: TranslationKey, + emphasizedValue: string, +): ReactNode { + const translated = t(key, { name: EMPHASIS_MARKER }); + const parts = translated.split(EMPHASIS_MARKER); + if (parts.length === 1) return translated; + + const nodes: ReactNode[] = []; + parts.forEach((part, index) => { + if (part.length > 0) nodes.push(part); + if (index < parts.length - 1) { + nodes.push({emphasizedValue}); + } + }); + return nodes; +} function getTourId(pathname: string): string | null { if (pathname.startsWith("/studio")) return "studio"; @@ -185,6 +209,7 @@ function NavItem({ } export function AppSidebar() { + const t = useT(); const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle(); const { pathname, search } = useRouterState({ select: (s) => ({ @@ -204,14 +229,8 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const [shutdownOpen, setShutdownOpen] = useState(false); - // Chat collapsible state — open by default, auto-expand on route entry const isChatRoute = pathname.startsWith("/chat"); const isStudioRoute = pathname === "/studio" || pathname.startsWith("/studio/"); - const [chatOpen, setChatOpen] = useState(true); - const [runsOpen, setRunsOpen] = useState(true); - - useEffect(() => { if (isChatRoute) setChatOpen(true); }, [isChatRoute]); - useEffect(() => { if (isStudioRoute) setRunsOpen(true); }, [isStudioRoute]); const scrollRef = useRef(null); const [scrolled, setScrolled] = useState(false); @@ -290,7 +309,7 @@ export function AppSidebar() { try { await renameChatItem(target.item, renameTrimmed); } catch (err) { - toast.error("Failed to rename chat", { + toast.error(translate("shell.toast.failedToRenameChat"), { description: err instanceof Error ? err.message : undefined, }); } @@ -300,7 +319,7 @@ export function AppSidebar() { const updated = await renameTrainingRun(target.run.id, nextRunDisplayName); emitTrainingRunUpdated(updated); } catch (err) { - toast.error("Failed to rename run", { + toast.error(translate("shell.toast.failedToRenameRun"), { description: err instanceof Error ? err.message : undefined, }); } @@ -320,14 +339,14 @@ export function AppSidebar() { try { await handleDeleteThread(target.item); } catch (err) { - toast.error("Failed to delete chat", { + toast.error(translate("shell.toast.failedToDeleteChat"), { description: err instanceof Error ? err.message : undefined, }); } return; } if (target.run.status === "running") { - toast.error("Cannot delete a running training run"); + toast.error(t("shell.toast.cannotDeleteRunningRun")); return; } try { @@ -337,7 +356,7 @@ export function AppSidebar() { } emitTrainingRunDeleted(target.run.id); } catch (err) { - toast.error("Failed to delete run", { + toast.error(translate("shell.toast.failedToDeleteRun"), { description: err instanceof Error ? err.message : undefined, }); } @@ -366,7 +385,7 @@ export function AppSidebar() { }); }} className="flex items-center gap-[6px] select-none" - aria-label="Unsloth home" + aria-label={t("shell.aria.home")} > - BETA + {t("shell.beta")} {!isMobile && ( @@ -387,7 +406,7 @@ export function AppSidebar() { type="button" onClick={togglePinned} className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Close sidebar" + aria-label={t("shell.aria.closeSidebar")} > @@ -397,7 +416,7 @@ export function AppSidebar() { sideOffset={6} className="tooltip-compact" > - Close sidebar + {t("shell.aria.closeSidebar")} )} @@ -412,7 +431,7 @@ export function AppSidebar() { type="button" onClick={togglePinned} className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Open sidebar" + aria-label={t("shell.aria.openSidebar")} > @@ -422,7 +441,7 @@ export function AppSidebar() { sideOffset={8} className="tooltip-compact" > - Open sidebar + {t("shell.aria.openSidebar")}
@@ -434,7 +453,7 @@ export function AppSidebar() { { @@ -446,7 +465,7 @@ export function AppSidebar() { /> i.id === search.compare)} disabled={chatDisabled} dataTour="chat-compare" @@ -459,7 +478,7 @@ export function AppSidebar() { /> { @@ -477,7 +496,7 @@ export function AppSidebar() { { @@ -489,7 +508,7 @@ export function AppSidebar() { { navigate({ to: "/data-recipes" }); @@ -499,7 +518,7 @@ export function AppSidebar() { { @@ -513,13 +532,16 @@ export function AppSidebar() { - {/* Recent Chats — hide on Studio only (Eyera fac13); chatOpen = ec695 clickability */} {!isStudioRoute && chatItems.length > 0 && ( - + - Recents + {t("shell.navigation.recents")} @@ -552,7 +574,7 @@ export function AppSidebar() { @@ -844,7 +872,9 @@ export function AppSidebar() { - {renamingTarget?.kind === "run" ? "Rename run" : "Rename chat"} + {renamingTarget?.kind === "run" + ? t("shell.dialog.renameRun.title") + : t("shell.dialog.renameChat.title")} @@ -868,14 +906,14 @@ export function AppSidebar() { variant="ghost" onClick={() => setRenamingTarget(null)} > - Cancel + {t("common.cancel")} diff --git a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx index f30af7fcf9..ed590f226e 100644 --- a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx +++ b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx @@ -5,6 +5,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { getAuthToken } from "@/features/auth"; +import { useT } from "@/i18n"; import { toastError, toastSuccess } from "@/shared/toast"; import { Camera } from "lucide-react"; import { useMemo, useRef, useState } from "react"; @@ -37,6 +38,7 @@ function readPersistedProfile(): { displayName: string; avatarDataUrl: string | } export function ProfilePersonalizationPanel() { + const t = useT(); const displayName = useUserProfileStore((s) => s.displayName); const avatarDataUrl = useUserProfileStore((s) => s.avatarDataUrl); const setDisplayName = useUserProfileStore((s) => s.setDisplayName); @@ -60,11 +62,11 @@ export function ProfilePersonalizationPanel() { setDisplayName(trimmed); const persisted = readPersistedProfile(); if (persisted && persisted.displayName === trimmed) { - toastSuccess("Profile name saved"); + toastSuccess(t("settings.profile.nameSaved")); } else { toastError( - "Could not persist profile name", - "Name updated for this session, but may not persist after reload.", + t("settings.profile.namePersistErrorTitle"), + t("settings.profile.namePersistErrorDescription"), ); } } @@ -78,17 +80,18 @@ export function ProfilePersonalizationPanel() { setAvatarDataUrl(dataUrl); const persisted = readPersistedProfile(); if (persisted && persisted.avatarDataUrl === dataUrl) { - toastSuccess("Profile photo updated"); + toastSuccess(t("settings.profile.photoUpdated")); } else { toastError( - "Could not persist profile photo", - "Photo updated for this session, but may not persist after reload.", + t("settings.profile.photoPersistErrorTitle"), + t("settings.profile.photoPersistErrorDescription"), ); } } catch (e) { - const message = e instanceof Error ? e.message : "Could not use this image."; + const message = + e instanceof Error ? e.message : t("settings.profile.imageUseError"); setImageError(message); - toastError("Could not update profile photo", message); + toastError(t("settings.profile.photoUpdateErrorTitle"), message); } }; @@ -115,7 +118,7 @@ export function ProfilePersonalizationPanel() { type="button" onClick={() => fileInputRef.current?.click()} className="absolute right-0 bottom-0 -translate-x-[15.625%] -translate-y-[15.625%] flex size-8 items-center justify-center rounded-full border border-border bg-background text-foreground shadow-sm transition-colors hover:bg-muted focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background" - aria-label="Change profile picture" + aria-label={t("settings.profile.changePicture")} > @@ -123,7 +126,7 @@ export function ProfilePersonalizationPanel() {
diff --git a/studio/frontend/src/features/settings/components/api-key-row.tsx b/studio/frontend/src/features/settings/components/api-key-row.tsx index ced6de17d1..9a6e1ee207 100644 --- a/studio/frontend/src/features/settings/components/api-key-row.tsx +++ b/studio/frontend/src/features/settings/components/api-key-row.tsx @@ -14,30 +14,39 @@ import { MoreHorizontalIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; +import { useT } from "@/i18n"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import type { ApiKey } from "../api/api-keys"; -function relative(iso: string | null): string { - if (!iso) return "never"; +type SettingsT = ReturnType; + +function relative(iso: string | null, t: SettingsT): string { + if (!iso) return t("settings.apiKeys.relativeNever"); const diff = Date.now() - new Date(iso).getTime(); const days = Math.floor(diff / 86400000); if (days < 1) { const hours = Math.floor(diff / 3600000); - if (hours < 1) return "just now"; - return `${hours}h ago`; + if (hours < 1) return t("settings.apiKeys.relativeJustNow"); + return t("settings.apiKeys.relativeHoursAgo", { count: hours }); } - if (days < 30) return `${days}d ago`; - if (days < 365) return `${Math.floor(days / 30)}mo ago`; - return `${Math.floor(days / 365)}y ago`; + if (days < 30) return t("settings.apiKeys.relativeDaysAgo", { count: days }); + if (days < 365) { + return t("settings.apiKeys.relativeMonthsAgo", { + count: Math.floor(days / 30), + }); + } + return t("settings.apiKeys.relativeYearsAgo", { + count: Math.floor(days / 365), + }); } -function expiresText(iso: string | null): string { - if (!iso) return "never"; +function expiresText(iso: string | null, t: SettingsT): string { + if (!iso) return t("settings.apiKeys.relativeNever"); const diff = new Date(iso).getTime() - Date.now(); - if (diff < 0) return "expired"; + if (diff < 0) return t("settings.apiKeys.expired"); const days = Math.floor(diff / 86400000); - if (days < 1) return "today"; - return `in ${days}d`; + if (days < 1) return t("settings.apiKeys.today"); + return t("settings.apiKeys.inDays", { count: days }); } export function ApiKeyRow({ @@ -47,6 +56,7 @@ export function ApiKeyRow({ apiKey: ApiKey; onRevoke: (key: ApiKey) => void; }) { + const t = useT(); const prefix = `sk-unsloth-${apiKey.key_prefix}…`; return (
@@ -64,11 +74,23 @@ export function ApiKeyRow({
- Created {relative(apiKey.created_at)} + + {t("settings.apiKeys.created", { + value: relative(apiKey.created_at, t), + })} + · - Used {relative(apiKey.last_used_at)} + + {t("settings.apiKeys.used", { + value: relative(apiKey.last_used_at, t), + })} + · - Expires {expiresText(apiKey.expires_at)} + + {t("settings.apiKeys.expires", { + value: expiresText(apiKey.expires_at, t), + })} +
@@ -77,7 +99,7 @@ export function ApiKeyRow({ variant="ghost" size="sm" className="size-7 p-0 opacity-0 transition-opacity group-hover:opacity-100 data-[state=open]:opacity-100 max-sm:!opacity-100 max-sm:size-9" - aria-label={`Actions for ${apiKey.name}`} + aria-label={t("settings.apiKeys.actionsFor", { name: apiKey.name })} > @@ -85,14 +107,14 @@ export function ApiKeyRow({ { await copyToClipboard(prefix); }}> - Copy prefix + {t("settings.apiKeys.copyPrefix")} onRevoke(apiKey)} className="text-destructive focus:text-destructive" > - Revoke token + {t("settings.apiKeys.revokeToken")} diff --git a/studio/frontend/src/features/settings/components/create-key-form.tsx b/studio/frontend/src/features/settings/components/create-key-form.tsx index a0f2d7f82d..93802bfad8 100644 --- a/studio/frontend/src/features/settings/components/create-key-form.tsx +++ b/studio/frontend/src/features/settings/components/create-key-form.tsx @@ -3,6 +3,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; +import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { useState } from "react"; import { createApiKey } from "../api/api-keys"; @@ -21,6 +22,7 @@ export function CreateKeyForm({ onCreated: (rawKey: string) => void; onError: (message: string) => void; }) { + const t = useT(); const [name, setName] = useState(""); const [expiry, setExpiry] = useState(null); const [loading, setLoading] = useState(false); @@ -33,8 +35,11 @@ export function CreateKeyForm({ const result = await createApiKey(name.trim(), expiry); onCreated(result.key); setName(""); - } catch (err) { - onError(err instanceof Error ? err.message : "Couldn't create access token."); + } catch { + // API helpers in ../api/api-keys.ts throw generic English Error + // messages; always use the translated message so zh-CN users do not + // see English text bleed through from internal exceptions. + onError(t("settings.apiKeys.createError")); } finally { setLoading(false); } @@ -49,9 +54,9 @@ export function CreateKeyForm({ setName(e.target.value)} - placeholder="Token name (e.g. production)" + placeholder={t("settings.apiKeys.tokenNamePlaceholder")} className="h-8 min-w-[180px] flex-1 text-sm" - aria-label="New access token name" + aria-label={t("settings.apiKeys.newAccessTokenName")} />
{EXPIRY_PRESETS.map((p) => { @@ -69,13 +74,15 @@ export function CreateKeyForm({ : "text-muted-foreground hover:text-foreground", )} > - {p.label} + {p.value === null ? t("settings.apiKeys.never") : p.label} ); })}
diff --git a/studio/frontend/src/features/settings/components/key-reveal-card.tsx b/studio/frontend/src/features/settings/components/key-reveal-card.tsx index 2b589e88fe..bdcb861c1d 100644 --- a/studio/frontend/src/features/settings/components/key-reveal-card.tsx +++ b/studio/frontend/src/features/settings/components/key-reveal-card.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { Button } from "@/components/ui/button"; +import { useT } from "@/i18n"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { cn } from "@/lib/utils"; import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; @@ -15,6 +16,7 @@ export function KeyRevealCard({ rawKey: string; onDone: () => void; }) { + const t = useT(); const [copied, setCopied] = useState(false); const handleCopy = async () => { @@ -32,7 +34,7 @@ export function KeyRevealCard({ className="size-3.5 text-emerald-600 dark:text-emerald-500" /> - New access token created + {t("settings.apiKeys.newTokenCreated")}

- Copy now — this won't be shown again. + {t("settings.apiKeys.copyNow")}

diff --git a/studio/frontend/src/features/settings/components/language-select.tsx b/studio/frontend/src/features/settings/components/language-select.tsx new file mode 100644 index 0000000000..9d30e06147 --- /dev/null +++ b/studio/frontend/src/features/settings/components/language-select.tsx @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { + LOCALES, + isSupportedLocale, + setLocale, + useT, + useLocale, +} from "@/i18n"; + +export function LanguageSelect() { + const t = useT(); + const locale = useLocale(); + + return ( + + ); +} diff --git a/studio/frontend/src/features/settings/components/theme-segmented.tsx b/studio/frontend/src/features/settings/components/theme-segmented.tsx index 1061995346..36e7062d8f 100644 --- a/studio/frontend/src/features/settings/components/theme-segmented.tsx +++ b/studio/frontend/src/features/settings/components/theme-segmented.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { cn } from "@/lib/utils"; +import { useT, type TranslationKey } from "@/i18n"; import { LaptopIcon, Moon02Icon, @@ -11,13 +12,18 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { motion, useReducedMotion } from "motion/react"; import { useTheme, type Theme } from "../stores/theme-store"; -const OPTIONS: { value: Theme; label: string; icon: typeof Sun02Icon }[] = [ - { value: "light", label: "Light", icon: Sun02Icon }, - { value: "dark", label: "Dark", icon: Moon02Icon }, - { value: "system", label: "System", icon: LaptopIcon }, +const OPTIONS: { + value: Theme; + labelKey: TranslationKey; + icon: typeof Sun02Icon; +}[] = [ + { value: "light", labelKey: "settings.appearance.theme.light", icon: Sun02Icon }, + { value: "dark", labelKey: "settings.appearance.theme.dark", icon: Moon02Icon }, + { value: "system", labelKey: "settings.appearance.theme.system", icon: LaptopIcon }, ]; export function ThemeSegmented() { + const t = useT(); const { theme, setTheme } = useTheme(); const reduced = useReducedMotion(); return ( @@ -49,7 +55,7 @@ export function ThemeSegmented() { /> )} - {opt.label} + {t(opt.labelKey)} ); })} diff --git a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx index e4cdccd2d7..90f66483b7 100644 --- a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx +++ b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -29,10 +30,13 @@ export type UpdateInstallSource = | "unknown"; type UpdateInstallSourceState = UpdateInstallSource | "loading"; -function getStudioUpdateInstructionLine(shell: UpdateShell): string { +function getStudioUpdateInstructionLine( + shell: UpdateShell, + t: ReturnType, +): string { return shell === "windows" - ? "Open PowerShell and run:" - : "Open Terminal and run:"; + ? t("settings.about.update.openPowerShell") + : t("settings.about.update.openTerminal"); } function isLocalInstallSource( @@ -59,6 +63,7 @@ function CopyableCommand({ command: string; copyLabel: string; }): ReactElement { + const t = useT(); const [copied, setCopied] = useState(false); const timerRef = useRef | null>(null); @@ -89,14 +94,26 @@ function CopyableCommand({ value={command} className="min-w-0 flex-1 bg-transparent px-2 py-1.5 font-mono text-[11px] text-foreground outline-none" title={command} - aria-label={`${copyLabel} text`} + aria-label={t("settings.about.update.commandText", { + label: copyLabel, + })} /> ); })} @@ -125,20 +131,20 @@ export function UsageExamples() { type="button" onClick={handleCopy} className="flex items-center gap-1 rounded px-1.5 py-1 text-[11px] text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Copy snippet" + aria-label={t("settings.apiKeys.copySnippet")} > - {copied ? "Copied" : "Copy"} + {copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")} 
           {snippets[lang]}
- Setup docs: + {t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( s.open); const activeTab = useSettingsDialogStore((s) => s.activeTab); const setActiveTab = useSettingsDialogStore((s) => s.setActiveTab); @@ -117,9 +133,9 @@ export function SettingsDialog() { "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", )} > - Settings + {t("settings.dialog.title")} - Manage your Unsloth Studio preferences. + {t("settings.dialog.description")}