ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-05-19 07:00:31 -07:00
commit 64dc11faa2
62 changed files with 6836 additions and 569 deletions
Download patch file Download diff file Expand all files Collapse all files

79
.github/workflows/lockfile-audit.yml vendored Normal file
View file

@ -0,0 +1,79 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# Fast, focused supply-chain audit of every checked-in lockfile.
#
# Runs scripts/lockfile_supply_chain_audit.py on PRs that touch any
# npm or cargo lockfile, on push to main, and on a daily schedule so
# newly-published IOCs surface even when no PR opens.
#
# Default behavior is "advisory": only public indicator-of-compromise
# strings, known-malicious pinned versions, and structurally broken
# lockfiles fail the build. Structural anomalies (missing integrity,
# non-default registry, etc.) are emitted as GitHub Actions warnings
# but do not block merges. This deliberately keeps the noise floor
# low while still failing the moment a checked-in lockfile starts
# pointing at known-bad bytes.
#
# This workflow is intentionally separate from security-audit.yml:
# - security-audit.yml is the umbrella job (pip-audit + npm audit +
# cargo audit + OSV + Semgrep + secret scanning + SBOM + ...);
# it takes ~25 minutes and runs only when dep manifests change.
# - lockfile-audit.yml is a ~30 second pure-Python parse + grep on
# the lockfiles themselves; it runs on every PR that even nudges
# a lockfile so reviewers always see the audit result inline.
name: Lockfile supply-chain audit
on:
pull_request:
paths:
- 'studio/frontend/package-lock.json'
- 'studio/backend/core/data_recipe/oxc-validator/package-lock.json'
- 'studio/package-lock.json'
- 'studio/src-tauri/Cargo.lock'
- 'scripts/lockfile_supply_chain_audit.py'
- '.github/workflows/lockfile-audit.yml'
push:
branches: [main]
paths:
- 'studio/frontend/package-lock.json'
- 'studio/backend/core/data_recipe/oxc-validator/package-lock.json'
- 'studio/package-lock.json'
- 'studio/src-tauri/Cargo.lock'
- 'scripts/lockfile_supply_chain_audit.py'
- '.github/workflows/lockfile-audit.yml'
schedule:
- cron: '37 5 * * *'
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
audit:
name: lockfile supply-chain audit
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- uses: actions/setup-python@v5
with:
python-version: '3.12'
- name: Verify audit script parses
run: python3 -c "import ast; ast.parse(open('scripts/lockfile_supply_chain_audit.py').read())"
- name: Run lockfile supply-chain audit
# Default mode: only known-malicious pinned versions, known IOC
# strings, and structurally broken lockfiles fail the build.
# Missing-integrity and other structural anomalies are emitted
# as ::warning:: annotations and do not gate merges.
run: python3 scripts/lockfile_supply_chain_audit.py

4
.gitignore vendored
View file

@ -229,5 +229,9 @@ log.txt
setup_leo.sh
server.pid
*.log
# Ignore stray lockfiles; real npm projects opt back in below (npm ci needs them).
package-lock.json
!studio/frontend/package-lock.json
!studio/backend/core/data_recipe/oxc-validator/package-lock.json
!studio/package-lock.json
llama.cpp/

2
.pre-commit-config.yaml
View file

@ -1,6 +1,6 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.12
rev: v0.15.13
hooks:
- id: ruff
args:

25
install.ps1
View file

@ -92,6 +92,7 @@ function Install-UnslothStudio {
$RepoRoot = ""
$TauriMode = $false
$SkipTorch = $false
$ShortcutsOnly = $false
$argList = $args
for ($i = 0; $i -lt $argList.Count; $i++) {
switch ($argList[$i]) {
@ -100,6 +101,7 @@ function Install-UnslothStudio {
"--no-torch" { $SkipTorch = $true }
"--verbose" { $script:UnslothVerbose = $true }
"-v" { $script:UnslothVerbose = $true }
"--shortcuts-only" { $ShortcutsOnly = $true }
"--package" {
$i++
if ($i -ge $argList.Count) {
@ -871,6 +873,19 @@ shell.Run cmd, 0, False
}
}
# Regen .lnk + launcher only; used by `unsloth studio update`.
if ($ShortcutsOnly) {
if ($TauriMode) { return }
$UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
if (-not (Test-Path -LiteralPath $UnslothExe)) {
Write-Host "[ERROR] unsloth.exe missing at $UnslothExe; run install.ps1 first." -ForegroundColor Red
# throw (not Exit-InstallFailure) so non-Tauri callers see rc != 0.
throw "unsloth.exe missing"
}
New-StudioShortcuts -UnslothExePath $UnslothExe
return
}
# ── Check winget ──
Write-TauriLog "STEP" "Checking system dependencies"
if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
@ -1285,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.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.4" unsloth-zoo }
if ($baseInstallExit -eq 0) {
$NoTorchReq = Find-NoTorchRuntimeFile
if ($NoTorchReq) {
@ -1293,7 +1308,7 @@ shell.Run cmd, 0, False
}
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.4" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -1331,7 +1346,7 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.4" unsloth-zoo }
if ($baseInstallExit -eq 0) {
$NoTorchReq = Find-NoTorchRuntimeFile
if ($NoTorchReq) {
@ -1339,7 +1354,7 @@ shell.Run cmd, 0, False
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.4" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -1367,7 +1382,7 @@ shell.Run cmd, 0, False
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.3" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.4" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)

26
install.sh
View file

@ -45,6 +45,7 @@ TAURI_MODE=false
_USER_PYTHON=""
_NO_TORCH_FLAG=false
_VERBOSE=false
_SHORTCUTS_ONLY=false
_next_is_package=false
_next_is_python=false
for arg in "$@"; do
@ -65,6 +66,7 @@ for arg in "$@"; do
--python) _next_is_python=true ;;
--no-torch) _NO_TORCH_FLAG=true ;;
--verbose|-v) _VERBOSE=true ;;
--shortcuts-only) _SHORTCUTS_ONLY=true ;;
esac
done
@ -1233,6 +1235,20 @@ elif grep -qi microsoft /proc/version 2>/dev/null; then
fi
step "platform" "$OS"
# Regen launcher/shortcuts only; used by `unsloth studio update`.
if [ "$_SHORTCUTS_ONLY" = true ]; then
# Tauri owns its own shortcuts.
if [ "$TAURI_MODE" != true ]; then
VENV_ABS_BIN="$VENV_DIR/bin"
if [ ! -x "$VENV_ABS_BIN/unsloth" ]; then
echo "ERROR: unsloth binary missing at '$VENV_ABS_BIN/unsloth'; run install.sh first." >&2
exit 1
fi
create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS"
fi
exit 0
fi
# ── Architecture detection & Python version ──
_ARCH=$(uname -m)
MAC_INTEL=false
@ -1849,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.3" unsloth-zoo
"unsloth>=2026.5.4" unsloth-zoo
_NO_TORCH_RT="$(_find_no_torch_runtime)"
if [ -n "$_NO_TORCH_RT" ]; then
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
@ -1857,7 +1873,7 @@ if [ "$_MIGRATED" = true ]; then
else
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.5.3" unsloth-zoo
"unsloth>=2026.5.4" unsloth-zoo
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2025,7 +2041,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.5.3" unsloth-zoo
"unsloth>=2026.5.4" unsloth-zoo
_NO_TORCH_RT="$(_find_no_torch_runtime)"
if [ -n "$_NO_TORCH_RT" ]; then
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
@ -2040,7 +2056,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.5.3" unsloth-zoo
--upgrade-package unsloth "unsloth>=2026.5.4" unsloth-zoo
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2072,7 +2088,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.3" --torch-backend=auto
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.4" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."

170
scripts/lockfile_supply_chain_audit.py
View file

@ -389,6 +389,18 @@ class Finding:
)
def _gha_escape(text: str) -> str:
"""Escape a string for use in a GitHub Actions `::warning::` /
`::error::` workflow command message. GH Actions truncates
annotation messages at the first newline unless `\\n` is
escaped as `%0A`; carriage returns and the percent sign need
matching escapes per the workflow-commands spec. Order matters:
`%` must be replaced first so the subsequent `%0A` / `%0D`
sequences are not double-encoded.
"""
return text.replace("%", "%25").replace("\r", "%0D").replace("\n", "%0A")
# ─────────────────────────────────────────────────────────────────────
# package-lock.json audit.
# ─────────────────────────────────────────────────────────────────────
@ -397,9 +409,35 @@ class Finding:
def audit_npm_lockfile(path: Path) -> list[Finding]:
findings: list[Finding] = []
if not path.exists():
# A missing requested lockfile is a config error, not a clean
# audit; surface it so a deleted default cannot pass silently.
findings.append(
Finding(
path = str(path),
package = "<root>",
kind = "missing-lockfile",
detail = (
"expected lockfile not found; refusing to silently "
"report a clean audit for a path that was not scanned"
),
)
)
return findings
raw = path.read_text(encoding = "utf-8")
try:
raw = path.read_text(encoding = "utf-8")
except OSError as exc:
# Permission denied, is-a-directory, broken-pipe etc. -- surface
# as a finding instead of crashing CI with a raw traceback.
findings.append(
Finding(
path = str(path),
package = "<root>",
kind = "unreadable-lockfile",
detail = f"could not read file: {exc}",
)
)
return findings
try:
lock = json.loads(raw)
except json.JSONDecodeError as exc:
@ -553,9 +591,32 @@ _PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$")
def audit_cargo_lockfile(path: Path) -> list[Finding]:
findings: list[Finding] = []
if not path.exists():
# See audit_npm_lockfile: missing lockfile is a finding.
findings.append(
Finding(
path = str(path),
package = "<root>",
kind = "missing-lockfile",
detail = (
"expected lockfile not found; refusing to silently "
"report a clean audit for a path that was not scanned"
),
)
)
return findings
raw = path.read_text(encoding = "utf-8")
try:
raw = path.read_text(encoding = "utf-8")
except OSError as exc:
findings.append(
Finding(
path = str(path),
package = "<root>",
kind = "unreadable-lockfile",
detail = f"could not read file: {exc}",
)
)
return findings
try:
import tomllib # type: ignore[import-not-found]
except ImportError:
@ -652,7 +713,31 @@ def audit_cargo_lockfile(path: Path) -> list[Finding]:
# ─────────────────────────────────────────────────────────────────────
DEFAULT_NPM_LOCKFILES = ("studio/frontend/package-lock.json",)
# Finding kinds split into BLOCKING vs ADVISORY for the default run mode.
# Blocking findings come from public supply-chain attack indicators (a
# version we know is malicious, a string an attacker would have to embed
# for an attack to work). Advisory findings are structural lockfile
# anomalies (missing integrity, non-default registry, etc.) -- they
# WARN the maintainer but do not block merges. Pass --strict to make
# every finding blocking (PR-5479-style behavior for opt-in adopters).
BLOCKING_KINDS: frozenset[str] = frozenset(
{
"blocked-known-malicious",
"known-ioc-string",
# Internal-failure kinds: a structurally broken lockfile MIGHT
# be hiding a real attack, so we keep these blocking too.
"malformed-lockfile",
"missing-lockfile",
"unreadable-lockfile",
"missing-toml-parser",
}
)
DEFAULT_NPM_LOCKFILES = (
"studio/frontend/package-lock.json",
"studio/backend/core/data_recipe/oxc-validator/package-lock.json",
"studio/package-lock.json",
)
DEFAULT_CARGO_LOCKFILES = ("studio/src-tauri/Cargo.lock",)
@ -671,7 +756,9 @@ def main(argv: list[str] | None = None) -> int:
default = None,
help = (
"Path to a package-lock.json (repeatable). "
"Default: studio/frontend/package-lock.json."
"Default: studio/frontend/package-lock.json, "
"studio/backend/core/data_recipe/oxc-validator/package-lock.json, "
"and studio/package-lock.json (Tauri CLI for desktop release)."
),
)
parser.add_argument(
@ -683,6 +770,18 @@ def main(argv: list[str] | None = None) -> int:
"Default: studio/src-tauri/Cargo.lock."
),
)
parser.add_argument(
"--strict",
action = "store_true",
help = (
"Treat every finding as blocking (exit 1). "
"Default mode only blocks on known-malicious versions, "
"indicator-of-compromise strings, or structurally broken "
"lockfiles; everything else is printed as an advisory "
"warning with exit 0. CI should use the default; local "
"audits aiming for zero noise can opt in via --strict."
),
)
args = parser.parse_args(argv)
# SF4: require a real justification (e.g. JIRA ticket id) for the
@ -715,8 +814,15 @@ def main(argv: list[str] | None = None) -> int:
return 0
root = Path(args.root).resolve()
npm_paths = [root / p for p in (args.npm_lockfile or DEFAULT_NPM_LOCKFILES)]
cargo_paths = [root / p for p in (args.cargo_lockfile or DEFAULT_CARGO_LOCKFILES)]
# Explicit --npm-lockfile/--cargo-lockfile scopes the scan to those
# paths; defaults apply only to the no-args CI invocation.
_user_explicit = args.npm_lockfile is not None or args.cargo_lockfile is not None
if _user_explicit:
npm_paths = [root / p for p in (args.npm_lockfile or ())]
cargo_paths = [root / p for p in (args.cargo_lockfile or ())]
else:
npm_paths = [root / p for p in DEFAULT_NPM_LOCKFILES]
cargo_paths = [root / p for p in DEFAULT_CARGO_LOCKFILES]
all_findings: list[Finding] = []
for p in npm_paths:
@ -734,17 +840,57 @@ def main(argv: list[str] | None = None) -> int:
)
return 0
# Split findings into blocking (known-malicious / IOC / structurally
# broken) and advisory (everything else, e.g. missing integrity on a
# registry-published tarball). In default mode advisory findings are
# printed but do not change the exit code; --strict treats every
# finding as blocking.
blocking = [f for f in all_findings if f.kind in BLOCKING_KINDS]
advisory = [f for f in all_findings if f.kind not in BLOCKING_KINDS]
if args.strict:
blocking = list(all_findings)
advisory = []
if advisory:
print(
f"\n[lockfile-audit] {len(advisory)} advisory finding(s) "
"(non-blocking; pass --strict to fail the build on these):\n",
file = sys.stderr,
)
for f in advisory:
# Surface in GitHub Actions UI as a warning annotation when run
# under Actions; harmless prefix elsewhere. GH Actions
# truncates annotation messages at the first newline unless
# newlines are escaped as `%0A`, so the full multi-line
# Finding (kind + path + package + detail) only renders in
# the UI after _gha_escape collapses it onto one line.
print(f"::warning::{_gha_escape(str(f))}", file = sys.stderr)
print(file = sys.stderr)
if not blocking:
print(
f"[lockfile-audit] OK: {len(advisory)} advisory finding(s), "
"0 blocking. Run with --strict to escalate advisory findings.",
flush = True,
)
return 0
print(
f"\n[lockfile-audit] FAIL: {len(all_findings)} finding(s):\n",
f"\n[lockfile-audit] FAIL: {len(blocking)} blocking finding(s):\n",
file = sys.stderr,
)
for f in all_findings:
print(str(f), file = sys.stderr)
for f in blocking:
# Same %-encoding rationale as the advisory branch above: the
# GH Actions annotation is truncated at the first newline
# unless the message is escaped.
print(f"::error::{_gha_escape(str(f))}", file = sys.stderr)
print(file = sys.stderr)
print(
"[lockfile-audit] Refusing to proceed. Each finding above is "
"either a structural lockfile anomaly or a public indicator-of-"
"compromise. Investigate before running `npm ci` or `cargo fetch`.",
"[lockfile-audit] Refusing to proceed. Each blocking finding "
"above is either a public indicator-of-compromise, a known-"
"malicious pinned version, or a structurally broken lockfile. "
"Investigate before running `npm ci` or `cargo fetch`.",
file = sys.stderr,
)
return 1

799
studio/backend/core/data_recipe/oxc-validator/package-lock.json generated Normal file
View file

@ -0,0 +1,799 @@
{
"name": "unsloth-oxc-validator-runtime",
"version": "0.0.1",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "unsloth-oxc-validator-runtime",
"version": "0.0.1",
"dependencies": {
"oxc-parser": "^0.123.0",
"oxlint": "^1.51.0"
}
},
"node_modules/@emnapi/core": {
"version": "1.10.0",
"resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.10.0.tgz",
"integrity": "sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==",
"license": "MIT",
"optional": true,
"peer": true,
"dependencies": {
"@emnapi/wasi-threads": "1.2.1",
"tslib": "^2.4.0"
}
},
"node_modules/@emnapi/runtime": {
"version": "1.10.0",
"resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.10.0.tgz",
"integrity": "sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==",
"license": "MIT",
"optional": true,
"peer": true,
"dependencies": {
"tslib": "^2.4.0"
}
},
"node_modules/@emnapi/wasi-threads": {
"version": "1.2.1",
"resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz",
"integrity": "sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==",
"license": "MIT",
"optional": true,
"peer": true,
"dependencies": {
"tslib": "^2.4.0"
}
},
"node_modules/@napi-rs/wasm-runtime": {
"version": "1.1.4",
"resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.4.tgz",
"integrity": "sha512-3NQNNgA1YSlJb/kMH1ildASP9HW7/7kYnRI2szWJaofaS1hWmbGI4H+d3+22aGzXXN9IJ+n+GiFVcGipJP18ow==",
"license": "MIT",
"optional": true,
"dependencies": {
"@tybys/wasm-util": "^0.10.1"
},
"funding": {
"type": "github",
"url": "https://github.com/sponsors/Brooooooklyn"
},
"peerDependencies": {
"@emnapi/core": "^1.7.1",
"@emnapi/runtime": "^1.7.1"
}
},
"node_modules/@oxc-parser/binding-android-arm-eabi": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-android-arm-eabi/-/binding-android-arm-eabi-0.123.0.tgz",
"integrity": "sha512-EHQ58z+6DbZWokMOKg5AB1KuwrXVgfbBLuuLFfzdc7bI5A4igvdvjKMhUv1VBV+0FABiUCOjNKUmMF7ugprwbQ==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-android-arm64": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-android-arm64/-/binding-android-arm64-0.123.0.tgz",
"integrity": "sha512-BK1E0zqNoHf38nTHjnGZ+olKHSKNHh65pChjY06yhaWYP8X7yNDqhQDA4neMPRqnPBgpN4/OW1oSMrdJgDi2aw==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-darwin-arm64": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-darwin-arm64/-/binding-darwin-arm64-0.123.0.tgz",
"integrity": "sha512-dkMPbtTbqU+cm+k4YGOBs4zAuq3Xu+wqjbGQvLAuVO7qHhNY4p5LBNudOmOoi0jxS8h1W6Jmlzv8MAKGpK+iDg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-darwin-x64": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-darwin-x64/-/binding-darwin-x64-0.123.0.tgz",
"integrity": "sha512-85pic0rCd59DGdM69jI9xE/Snb2KtrfiU48QigjJXjzxUOenGvH4SAFIjFpO/2ZnI3Kz50D8pht4jKN3t2022Q==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-freebsd-x64": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-freebsd-x64/-/binding-freebsd-x64-0.123.0.tgz",
"integrity": "sha512-mjEiW6z7JtaiHMK/8aJic1lfjkKpzFwK2XFNmm187BFbtDamjGVuKNr2TEyrFEYJyZc217wokR1wrYeZGBQo4Q==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"freebsd"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-arm-gnueabihf": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-0.123.0.tgz",
"integrity": "sha512-mYxigPtGt6SZfhNZBIJfuDM92cLo8XUW08WuKxzHvcmWu6xndLqwLp99Vg4uHke1AXicQEHU3Wri2X9bHF0Vlw==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-arm-musleabihf": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-0.123.0.tgz",
"integrity": "sha512-ttWirDC9eUBn0R4Tzz3aeDaLrx9drPdNiLJ8MXeDBFxd6cwLfTIC27qjsdfGpn942tkVIZY3sjWAnvbwDDjX7g==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-arm64-gnu": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-0.123.0.tgz",
"integrity": "sha512-apAHyoMNRYT+2G98Y14caZmsr5LD9PsWpGI7nXmSwK26LGiQneCU6HvHQ+d+AX+RJ5TTWZtEb2RD7OLqAC0cYQ==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-arm64-musl": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm64-musl/-/binding-linux-arm64-musl-0.123.0.tgz",
"integrity": "sha512-3r99Qa4egjO/iXUBxTlN6Ddt1YkLifG6olzvj8gkoKEK2U/MOW7mQfXRyBmuoMgmZ7O4vk41gO3d21c6VcN3yQ==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-ppc64-gnu": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-0.123.0.tgz",
"integrity": "sha512-Hr/Z24kUE4pjJs346g80WDwjyJGrxiw6hExJuOiME/76ZFz68y5L11UzprRkW9FN4HxBB7tLZ/fytczV2fEsiA==",
"cpu": [
"ppc64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-riscv64-gnu": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-0.123.0.tgz",
"integrity": "sha512-sxjbhs+8WXeuoLnZ2rBmQ96gPdq3SCmz24reIltsKLUt1EDMgdaQsr7RqwBphw3QAImkMtlPQfAWDWwZyo0xDg==",
"cpu": [
"riscv64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-riscv64-musl": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-0.123.0.tgz",
"integrity": "sha512-d6xHHhqldA/W+VC7v8uHs24zM69Ad3HnHQ45h+uuBhCsbZx3d0E0wL2K3uJ5mYKTR6UPMFk9VMXcHWwvg1PRZQ==",
"cpu": [
"riscv64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-s390x-gnu": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-0.123.0.tgz",
"integrity": "sha512-+di9A5wJQlv0VodyhADjJ2rC4geyHY+uhJDl3TFjMgYhhlgLZchi9uHD5mfiUEDWHt1x7/eU2u1ge3LLazZmFw==",
"cpu": [
"s390x"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-x64-gnu": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-x64-gnu/-/binding-linux-x64-gnu-0.123.0.tgz",
"integrity": "sha512-sh7pw2g/u6LE1TaRRQsV9Kv9+1y+CywaaNwWWP+3bnEPk/L692oTG0hmEviUlawI8v3OGC+AhbjtAD+HXWQAkg==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-linux-x64-musl": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-x64-musl/-/binding-linux-x64-musl-0.123.0.tgz",
"integrity": "sha512-S+LoD8PiJ639JwIqK1knIeqAyYkeCbLHtAgfapszKX0yVCaYP+aer8dJxL25de9qcDjvYWVrYCkuDZzHmOl2Xw==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-openharmony-arm64": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-openharmony-arm64/-/binding-openharmony-arm64-0.123.0.tgz",
"integrity": "sha512-/65vryK11q1I+k+7ukDlwZOxUFCLYsoZBZPGZHyet5bIP5e3D8mV3uCuvpWZ9Hoe6vUZFw/nAfCrX59MeuJPgw==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"openharmony"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-wasm32-wasi": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-wasm32-wasi/-/binding-wasm32-wasi-0.123.0.tgz",
"integrity": "sha512-y4OsMGQiAbZzj2Rq0LEfvhR48rQDvbvqsl/dPdn4tdf+z3H79nZuR+lQ/+KUGjD30vpVGem138sBWHFj9UR+Vg==",
"cpu": [
"wasm32"
],
"license": "MIT",
"optional": true,
"dependencies": {
"@napi-rs/wasm-runtime": "^1.1.2"
},
"engines": {
"node": ">=14.0.0"
}
},
"node_modules/@oxc-parser/binding-win32-arm64-msvc": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-0.123.0.tgz",
"integrity": "sha512-9lBqI6AXAkjYavkdpizNU3Q51uoVYfp9FJPx19hnCEdPku1jSgzSnvgmCvhCue0GziIvIvIdWgZ41wXQ3EOoBw==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-win32-ia32-msvc": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-ia32-msvc/-/binding-win32-ia32-msvc-0.123.0.tgz",
"integrity": "sha512-zJbqBHwSUB7CyvAONy9ewGtQwcQj+ylOhYGETvUPp3KIYx7lolj4Gayof7iA22SU5eMSjO5COL0c8wYhmn9agA==",
"cpu": [
"ia32"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-parser/binding-win32-x64-msvc": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-x64-msvc/-/binding-win32-x64-msvc-0.123.0.tgz",
"integrity": "sha512-q7RZvglQvGo3RX5ljtcGSabu2B2c0oDU/6xC3sBMhsV5KRo0PvyxLdordbEN31NTfuZu4Sgl86C76cAURZIHWA==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxc-project/types": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.123.0.tgz",
"integrity": "sha512-YtECP/y8Mj1lSHiUWGSRzy/C6teUKlS87dEfuVKT09LgQbUsBW1rNg+MiJ4buGu3yuADV60gbIvo9/HplA56Ew==",
"license": "MIT",
"funding": {
"url": "https://github.com/sponsors/Boshen"
}
},
"node_modules/@oxlint/binding-android-arm-eabi": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-android-arm-eabi/-/binding-android-arm-eabi-1.64.0.tgz",
"integrity": "sha512-2r6Nq3XXGLHEXKkSj8JtmJ6N4gDw431DPFOg0ZoJHlNjnG6HVMm/ksQ10m0HJ8WBvwgMe1L50UHPaYZutCRPCw==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-android-arm64": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-android-arm64/-/binding-android-arm64-1.64.0.tgz",
"integrity": "sha512-ePJMpePgg7fBv+L/hVx1xXRU5/5gd5m0obLA6hPEfLXF3GjpR8idIDbY1dhQYhyz1ms2wdTccSboo6KEd2Oxtg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-darwin-arm64": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-darwin-arm64/-/binding-darwin-arm64-1.64.0.tgz",
"integrity": "sha512-U4DMLQd10gJLuoSTLSGbfv3bGjTlUNsScm9Dgb8wwBqmCzidf1pE1pXV4doGNxqwH3KtVng1AGTINA0NvkGLvQ==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-darwin-x64": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-darwin-x64/-/binding-darwin-x64-1.64.0.tgz",
"integrity": "sha512-GoRIL48QWm4/TAvjN8pB1nAG+1/uqc9EdnWT9zqHeb6wsmjZtywj8VRe5aGW47Fdb64YtLOsdLqVxOvQuz98Wg==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-freebsd-x64": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-freebsd-x64/-/binding-freebsd-x64-1.64.0.tgz",
"integrity": "sha512-5dFkv4tkg7PxJJGS9/OjrJwjhuHczrd3OQOkRE0wHcLM+ncUnULtzEPWjqGOxTXxZnLWcB91bGiIznx89TVXyQ==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"freebsd"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-arm-gnueabihf": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.64.0.tgz",
"integrity": "sha512-jsBqMLl/uOL5+Kq/+BtK9FrmiNGUbx8SiyZXv+WlUxA45KuwcLu9BfiSIL3I3DBDgWM3yZizDITnTK9BcqNBQg==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-arm-musleabihf": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-1.64.0.tgz",
"integrity": "sha512-1lrj8At/Uuc9GhjrVFBQo0NEjfBrTkzpmtHIGAhNnIXqn1CAyGL+qrztUsXb2GIluJrpl9Q7qRLJOb/NqydacQ==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-arm64-gnu": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.64.0.tgz",
"integrity": "sha512-HpSQbubwh03mMhAdy2BYtad/fsY8vDFHDAb6bUwuCYg2VD3xCQgn6ArKcO0oZyLCheacKTv4PrF3Mfu5hgoE2g==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-arm64-musl": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.64.0.tgz",
"integrity": "sha512-00QQ0h0Y7u0G69BgiH3+ky2aaq/QvkDL6DYok8htIuJHxybiux5aQ8jwmg8qIk9wha6UagUP2BAwAzbemcJbpg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-ppc64-gnu": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.64.0.tgz",
"integrity": "sha512-2GaimTV6EMW+s5HS0An3oGbQme3BgHswvfVdGk3EB57Xe9+/gyT+Qd7lNVzb3rtir52vbIPzXfaYArzs5b5zcw==",
"cpu": [
"ppc64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-riscv64-gnu": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-1.64.0.tgz",
"integrity": "sha512-H46AtFb9wypjoVwGdlxrm0DsD809NGmtiK9HiyPKTxkSte2YjhC4S+00rOIrwCaxcyPiGid3Y3OMXp5KMAkGZw==",
"cpu": [
"riscv64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-riscv64-musl": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-1.64.0.tgz",
"integrity": "sha512-HEgsidjjvvyzdg82icYkuFCf7REDV7B9JFwbIMbVwrKLBY0MrXX+bku3POn/hduZ2yW91IyVDUMq0Bf02KwXQw==",
"cpu": [
"riscv64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-s390x-gnu": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.64.0.tgz",
"integrity": "sha512-Axvm8qryotmKN00P5w4JapaSjvP2LOSbdbBJiX+2SuHd3QzhW7TUc8skqgw+ahQZ5DmzEYeHCqauvW8f32Ns6Q==",
"cpu": [
"s390x"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-x64-gnu": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.64.0.tgz",
"integrity": "sha512-cR60vSd7+m+KRZ3GQGfDxWwahW5RMXg0qlGvAluZr0fTUYvw0H9N9AXAF/M/PMqgytyqvVNmBAkJG9l7U30Y1g==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-linux-x64-musl": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-linux-x64-musl/-/binding-linux-x64-musl-1.64.0.tgz",
"integrity": "sha512-2u/aPZ9pEg7HnvZPDsHxUGNnrpr4qaHi+mCgLgpt+LYRzPrS4Px4wPfkIdRdr2GvKnaYyt+XSlto0Vm5sbStTg==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-openharmony-arm64": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-openharmony-arm64/-/binding-openharmony-arm64-1.64.0.tgz",
"integrity": "sha512-kfhkGfCdoXLSxEkrhDlJrvBYajGmq+ma4EMc53dsOWTq+rIBOlI0vTBmpZNnM5oH2LY/K/w1HAK+UQEgjgpVUg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"openharmony"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-win32-arm64-msvc": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.64.0.tgz",
"integrity": "sha512-r/cNKBFieONoVu2bb1KkVouq9W+edDUgHumXJGphCRRj+U0xaD4nanrw8ZOqo0IsutPkEM4vCcGBpak6x5aXMg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-win32-ia32-msvc": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-win32-ia32-msvc/-/binding-win32-ia32-msvc-1.64.0.tgz",
"integrity": "sha512-tUw0xUUwEFVZbpJoeCblkv8SJA4Xz3CdXCJbAnBsiNLyxDrk2tLcxEAS6M73Q7hHHDg3OtwI8vZVK3t5RJt4Gw==",
"cpu": [
"ia32"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@oxlint/binding-win32-x64-msvc": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/@oxlint/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.64.0.tgz",
"integrity": "sha512-9CBR+LO0JVST87fNTzzNxS5I29jIUO5gxT9i9+M3SDHHALElj9sY1Prf12tad3vIRC6OD7Ehtvvh+sn13vSwHw==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": "^20.19.0 || >=22.12.0"
}
},
"node_modules/@tybys/wasm-util": {
"version": "0.10.2",
"resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.2.tgz",
"integrity": "sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg==",
"license": "MIT",
"optional": true,
"dependencies": {
"tslib": "^2.4.0"
}
},
"node_modules/oxc-parser": {
"version": "0.123.0",
"resolved": "https://registry.npmjs.org/oxc-parser/-/oxc-parser-0.123.0.tgz",
"integrity": "sha512-F6ak0tFc01ZGbl5KxvLDQ2K005Z086mp3ByCQBDhUjqXLkapGUkMuJSsYixncdEpkLlcRDcruHR71LD339ADUA==",
"license": "MIT",
"dependencies": {
"@oxc-project/types": "^0.123.0"
},
"engines": {
"node": "^20.19.0 || >=22.12.0"
},
"funding": {
"url": "https://github.com/sponsors/Boshen"
},
"optionalDependencies": {
"@oxc-parser/binding-android-arm-eabi": "0.123.0",
"@oxc-parser/binding-android-arm64": "0.123.0",
"@oxc-parser/binding-darwin-arm64": "0.123.0",
"@oxc-parser/binding-darwin-x64": "0.123.0",
"@oxc-parser/binding-freebsd-x64": "0.123.0",
"@oxc-parser/binding-linux-arm-gnueabihf": "0.123.0",
"@oxc-parser/binding-linux-arm-musleabihf": "0.123.0",
"@oxc-parser/binding-linux-arm64-gnu": "0.123.0",
"@oxc-parser/binding-linux-arm64-musl": "0.123.0",
"@oxc-parser/binding-linux-ppc64-gnu": "0.123.0",
"@oxc-parser/binding-linux-riscv64-gnu": "0.123.0",
"@oxc-parser/binding-linux-riscv64-musl": "0.123.0",
"@oxc-parser/binding-linux-s390x-gnu": "0.123.0",
"@oxc-parser/binding-linux-x64-gnu": "0.123.0",
"@oxc-parser/binding-linux-x64-musl": "0.123.0",
"@oxc-parser/binding-openharmony-arm64": "0.123.0",
"@oxc-parser/binding-wasm32-wasi": "0.123.0",
"@oxc-parser/binding-win32-arm64-msvc": "0.123.0",
"@oxc-parser/binding-win32-ia32-msvc": "0.123.0",
"@oxc-parser/binding-win32-x64-msvc": "0.123.0"
}
},
"node_modules/oxlint": {
"version": "1.64.0",
"resolved": "https://registry.npmjs.org/oxlint/-/oxlint-1.64.0.tgz",
"integrity": "sha512-Star3SNpWPeWFPw7kRXIhXUSn6fdiAl25q15CQzH/9WaOtG6e9CWTc25vNZOCr4PE1yEP1GtKJKIKglhj3OmEQ==",
"license": "MIT",
"bin": {
"oxlint": "bin/oxlint"
},
"engines": {
"node": "^20.19.0 || >=22.12.0"
},
"funding": {
"url": "https://github.com/sponsors/Boshen"
},
"optionalDependencies": {
"@oxlint/binding-android-arm-eabi": "1.64.0",
"@oxlint/binding-android-arm64": "1.64.0",
"@oxlint/binding-darwin-arm64": "1.64.0",
"@oxlint/binding-darwin-x64": "1.64.0",
"@oxlint/binding-freebsd-x64": "1.64.0",
"@oxlint/binding-linux-arm-gnueabihf": "1.64.0",
"@oxlint/binding-linux-arm-musleabihf": "1.64.0",
"@oxlint/binding-linux-arm64-gnu": "1.64.0",
"@oxlint/binding-linux-arm64-musl": "1.64.0",
"@oxlint/binding-linux-ppc64-gnu": "1.64.0",
"@oxlint/binding-linux-riscv64-gnu": "1.64.0",
"@oxlint/binding-linux-riscv64-musl": "1.64.0",
"@oxlint/binding-linux-s390x-gnu": "1.64.0",
"@oxlint/binding-linux-x64-gnu": "1.64.0",
"@oxlint/binding-linux-x64-musl": "1.64.0",
"@oxlint/binding-openharmony-arm64": "1.64.0",
"@oxlint/binding-win32-arm64-msvc": "1.64.0",
"@oxlint/binding-win32-ia32-msvc": "1.64.0",
"@oxlint/binding-win32-x64-msvc": "1.64.0"
},
"peerDependencies": {
"oxlint-tsgolint": ">=0.22.1"
},
"peerDependenciesMeta": {
"oxlint-tsgolint": {
"optional": true
}
}
},
"node_modules/tslib": {
"version": "2.8.1",
"resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz",
"integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==",
"license": "0BSD",
"optional": true
}
}
}

60
studio/backend/core/inference/chat_template_helpers.py Normal file
View file

@ -0,0 +1,60 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Dependency-light wrapper around tokenizer.apply_chat_template with a
kwarg fallback for templates that reject reasoning/tools args.
"""
from typing import Optional
def apply_chat_template_for_generation(
tokenizer,
messages: list,
*,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> str:
"""Render the chat prompt. Try richest kwargs first; drop one
group at a time on TypeError. Jinja / missing-variable errors
propagate."""
reasoning_kwargs: dict = {}
if enable_thinking is not None:
reasoning_kwargs["enable_thinking"] = enable_thinking
if reasoning_effort is not None:
reasoning_kwargs["reasoning_effort"] = reasoning_effort
if preserve_thinking is not None:
reasoning_kwargs["preserve_thinking"] = preserve_thinking
attempts: list[dict] = []
if tools and reasoning_kwargs:
attempts.append({"tools": tools, **reasoning_kwargs})
if tools:
attempts.append({"tools": tools})
if reasoning_kwargs:
attempts.append(dict(reasoning_kwargs))
attempts.append({})
last_exc: Optional[Exception] = None
for kwargs in attempts:
try:
return tokenizer.apply_chat_template(
messages,
tokenize = False,
add_generation_prompt = True,
**kwargs,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError(
"apply_chat_template_for_generation: no attempt produced a result"
)

138
studio/backend/core/inference/inference.py
View file

@ -839,6 +839,74 @@ class InferenceBackend:
cancel_event = cancel_event, _adapter_state = use_adapter, **gen_kwargs
)
def generate_chat_completion_with_tools(
self,
messages: list,
tools: list,
system_prompt: str = "",
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_new_tokens: int = 2048,
repetition_penalty: float = 1.0,
cancel_event = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
):
"""Run an agentic tool loop on top of ``generate_chat_response``.
Yields the same event-dict protocol used by the GGUF path so
the route layer can stream both backends through one helper.
Each event is one of:
* ``{"type": "status", "text": ...}``
* ``{"type": "content", "text": cumulative_text}``
* ``{"type": "tool_start", "tool_name", "tool_call_id", "arguments"}``
* ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}``
"""
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from core.inference.tools import execute_tool
def _single_turn(conv: list):
# conv already has the system message -- avoid double-prepend.
yield from self._generate_chat_response_inner(
messages = conv,
system_prompt = "",
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
initial = list(messages)
if system_prompt:
initial = [{"role": "system", "content": system_prompt}] + initial
yield from run_safetensors_tool_loop(
single_turn = _single_turn,
messages = initial,
tools = tools,
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
)
def generate_chat_response(
self,
messages: list,
@ -851,10 +919,20 @@ class InferenceBackend:
max_new_tokens: int = 256,
repetition_penalty: float = 1.0,
cancel_event = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""
Generate response for text or vision models.
The generation lock is acquired by the background generation thread.
``tools`` / ``enable_thinking`` / ``reasoning_effort`` /
``preserve_thinking`` are forwarded into
``tokenizer.apply_chat_template`` so templates that understand
these kwargs (Qwen3, Llama 3.1+, gpt-oss harmony, ...) advertise
the tool schemas and reasoning controls to the model.
"""
yield from self._generate_chat_response_inner(
messages = messages,
@ -867,6 +945,10 @@ class InferenceBackend:
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
def _generate_chat_response_inner(
@ -882,6 +964,10 @@ class InferenceBackend:
repetition_penalty: float = 1.0,
cancel_event = None,
_adapter_state = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""
Inner generation logic. Called by both generate_chat_response
@ -981,8 +1067,13 @@ class InferenceBackend:
f"Please use a model that includes a chat template, or manually set "
f"one via tokenizer.chat_template before inference."
)
formatted_prompt = tokenizer.apply_chat_template(
template_messages, tokenize = False, add_generation_prompt = True
formatted_prompt = self._apply_chat_template_for_generation(
tokenizer,
template_messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
@ -1319,20 +1410,9 @@ class InferenceBackend:
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
"""Check if the given (or active) model uses the gpt-oss harmony protocol."""
name = (model_name or self.active_model_name or "").lower()
try:
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER
from utils.datasets import is_gpt_oss_model_name
# Exact match
if MODEL_TO_TEMPLATE_MAPPER.get(name) == "gpt-oss":
return True
# Partial match (e.g. name-bnb-4bit variants)
for key, tmpl in MODEL_TO_TEMPLATE_MAPPER.items():
if tmpl == "gpt-oss" and (key in name or name in key):
return True
except Exception:
pass
return "gpt-oss" in name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
def generate_stream(
self,
@ -1715,6 +1795,34 @@ class InferenceBackend:
"Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS"
)
def _apply_chat_template_for_generation(
self,
tokenizer,
messages: list,
*,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> str:
"""Render the chat prompt, peeling kwargs the template does not
understand. Delegates to the dependency-light helper module so
the fallback chain can be unit-tested without pulling unsloth /
torch into the test sandbox.
"""
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
)
return apply_chat_template_for_generation(
tokenizer,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str:
if not self.active_model_name or self.active_model_name not in self.models:
logger.error("No active model available")

825
studio/backend/core/inference/llama_cpp.py
View file

File diff suppressed because it is too large Load diff

2
studio/backend/core/inference/llama_server_args.py
View file

@ -148,6 +148,8 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
# MTP path (llama.cpp #22673).
"--spec-draft-n-max",
"--spec-draft-n-min",
"--spec-draft-p-min",
"--spec-draft-p-split",
"--spec-ngram-mod-n-match",
"--spec-ngram-mod-n-min",
"--spec-ngram-mod-n-max",

113
studio/backend/core/inference/mlx_inference.py
View file

@ -157,10 +157,59 @@ class MLXInferenceBackend:
"audio_type": None,
"has_audio_input": False,
}
# Capture chat_template_info so the worker IPC reply can ship
# it back to the parent and the route layer classifies
# capabilities the same way as the transformers / GGUF paths.
self._populate_chat_template_info(model_name)
logger.info("Model %s loaded successfully", model_name)
return True
def _populate_chat_template_info(self, model_name: str) -> None:
"""Mirror InferenceBackend._load_chat_template_info for MLX.
Stores ``chat_template_info`` on ``self.models[model_name]``
with the resolved ``tokenizer.chat_template`` so
``_detect_safetensors_features`` (route layer) sees the same
template the model actually uses."""
entry = self.models.get(model_name)
if not entry:
return
tok = entry.get("tokenizer")
if tok is None:
proc = entry.get("processor")
tok = getattr(proc, "tokenizer", None) if proc else None
info = {
"has_template": False,
"template": None,
"format_type": "generic",
"special_tokens": {},
"template_name": None,
}
try:
tpl = getattr(tok, "chat_template", None)
if tpl:
info["has_template"] = True
info["template"] = tpl
lower = tpl.lower()
if "start_header_id" in lower and "end_header_id" in lower:
info["format_type"] = "llama3"
elif "[inst]" in lower and "[/inst]" in lower:
info["format_type"] = "mistral"
elif "<|im_start|>" in lower and "<|im_end|>" in lower:
info["format_type"] = "chatml"
else:
info["format_type"] = "custom"
special = {}
for attr in ("bos_token", "eos_token", "pad_token"):
val = getattr(tok, attr, None)
if val:
special[attr] = val
info["special_tokens"] = special
except Exception as exc:
logger.warning("MLX chat_template_info capture failed: %s", exc)
entry["chat_template_info"] = info
def unload_model(self, model_name: str) -> bool:
import mlx.core as mx
import gc
@ -197,6 +246,14 @@ class MLXInferenceBackend:
max_new_tokens = 256,
repetition_penalty = 1.0,
cancel_event = None,
# Reasoning / tool kwargs forwarded by the route + worker -- the
# MLX path renders the template via apply_chat_template_for_
# generation so these are honoured the same way as the
# transformers path.
tools = None,
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
) -> Generator[str, None, None]:
if self._model is None:
raise RuntimeError("No model loaded")
@ -239,6 +296,10 @@ class MLXInferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
else:
yield from self._generate_text(
@ -250,6 +311,10 @@ class MLXInferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
def _generate_text(
@ -262,14 +327,26 @@ class MLXInferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event,
*,
tools = None,
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
):
from mlx_lm import stream_generate
from mlx_lm.sample_utils import make_sampler, make_logits_processors
prompt = self._tokenizer.apply_chat_template(
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
)
prompt = apply_chat_template_for_generation(
self._tokenizer,
messages,
tokenize = False,
add_generation_prompt = True,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
if prompt is None:
raise RuntimeError(
@ -343,20 +420,38 @@ class MLXInferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event,
*,
tools = None,
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
):
from mlx_vlm import stream_generate as vlm_stream
# Apply chat template
chat_fn = getattr(self._processor, "apply_chat_template", None)
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
)
# Pick the chat-template-aware caller: processors that expose
# their own apply_chat_template + chat_template attr (e.g.
# Qwen2.5-VL) use it directly; otherwise fall back to the
# nested tokenizer.
chat_target = self._processor
if (
chat_fn is None
getattr(self._processor, "apply_chat_template", None) is None
or not hasattr(self._processor, "chat_template")
or self._processor.chat_template is None
):
tok = getattr(self._processor, "tokenizer", self._processor)
chat_fn = tok.apply_chat_template
chat_target = getattr(self._processor, "tokenizer", self._processor)
prompt = chat_fn(messages, tokenize = False, add_generation_prompt = True)
prompt = apply_chat_template_for_generation(
chat_target,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# For VLM: always use mlx_vlm's stream_generate which handles
# pixel_values properly (passes None for text-only, image for VLM)

142
studio/backend/core/inference/orchestrator.py
View file

@ -449,6 +449,10 @@ class InferenceOrchestrator:
repetition_penalty: float = 1.0,
cancel_event = None,
use_adapter = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""Dispatched generation — sends command without holding _gen_lock.
@ -494,6 +498,14 @@ class InferenceOrchestrator:
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
cmd["enable_thinking"] = enable_thinking
if reasoning_effort is not None:
cmd["reasoning_effort"] = reasoning_effort
if preserve_thinking is not None:
cmd["preserve_thinking"] = preserve_thinking
# Create mailbox BEFORE sending command
mailbox: queue.Queue = queue.Queue()
@ -695,6 +707,13 @@ class InferenceOrchestrator:
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
}
# Mirror chat_template_info so routes can classify
# capabilities without re-entering the subprocess.
_tpl_info = model_info.get("chat_template_info")
if isinstance(_tpl_info, dict):
self.models[self.active_model_name]["chat_template_info"] = (
_tpl_info
)
self.loading_models.discard(model_name)
logger.info(
"Model '%s' loaded successfully in subprocess", model_name
@ -770,8 +789,18 @@ class InferenceOrchestrator:
max_new_tokens: int = 256,
repetition_penalty: float = 1.0,
cancel_event = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""Generate response, streaming tokens from subprocess."""
"""Generate response, streaming tokens from subprocess.
Optional ``tools`` / ``enable_thinking`` / ``reasoning_effort`` /
``preserve_thinking`` kwargs are forwarded into the worker so
``tokenizer.apply_chat_template`` can render tool schemas and
reasoning controls when the template understands them.
"""
yield from self._generate_inner(
messages = messages,
system_prompt = system_prompt,
@ -784,6 +813,88 @@ class InferenceOrchestrator:
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
use_adapter = None,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
def generate_chat_completion_with_tools(
self,
messages: list,
tools: list,
system_prompt: str = "",
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_tokens: Optional[int] = None,
repetition_penalty: float = 1.0,
cancel_event = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
use_adapter: Optional[Union[bool, str]] = None,
**_unused,
):
"""Run the safetensors agentic tool loop in this (parent)
process, calling the worker for each generation turn.
Yields the same event dicts as the GGUF tool loop so the route
layer can stream both backends through one helper. See
``safetensors_agentic.run_safetensors_tool_loop`` for the
event protocol.
"""
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from core.inference.tools import execute_tool
max_new_tokens = max_tokens if max_tokens and max_tokens > 0 else 2048
def _single_turn(conv: list):
# ``conv`` already carries any system message because the
# loop appends to a list seeded with system+user above.
common_kwargs = dict(
messages = conv,
system_prompt = "",
image = None,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
if use_adapter is not None:
yield from self.generate_with_adapter_control(
use_adapter = use_adapter,
**common_kwargs,
)
else:
yield from self.generate_chat_response(**common_kwargs)
initial = list(messages)
if system_prompt:
initial = [{"role": "system", "content": system_prompt}] + initial
yield from run_safetensors_tool_loop(
single_turn = _single_turn,
messages = initial,
tools = tools,
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
)
def generate_with_adapter_control(
@ -817,6 +928,10 @@ class InferenceOrchestrator:
repetition_penalty: float = 1.0,
cancel_event = None,
use_adapter = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""Inner generation logic — sends command to subprocess, yields tokens.
@ -853,6 +968,10 @@ class InferenceOrchestrator:
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
def _generate_locked(
@ -868,6 +987,10 @@ class InferenceOrchestrator:
repetition_penalty: float = 1.0,
cancel_event = None,
use_adapter = None,
tools: Optional[list] = None,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""Actual generation logic — must be called under _gen_lock."""
request_id = str(uuid.uuid4())
@ -893,6 +1016,16 @@ class InferenceOrchestrator:
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
# Only forward template kwargs the caller actually set so older
# workers that ignore unknown keys still work.
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
cmd["enable_thinking"] = enable_thinking
if reasoning_effort is not None:
cmd["reasoning_effort"] = reasoning_effort
if preserve_thinking is not None:
cmd["preserve_thinking"] = preserve_thinking
try:
self._send_cmd(cmd)
@ -1200,6 +1333,13 @@ class InferenceOrchestrator:
return self.models[self.active_model_name].get("is_vision", False)
return False
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
"""Parent-side gpt-oss detection so the safetensors route can run
the same guard without an IPC round-trip to the subprocess."""
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
# ========== GLOBAL INSTANCE ==========
_inference_backend = None

392
studio/backend/core/inference/safetensors_agentic.py Normal file
View file

@ -0,0 +1,392 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Safetensors/transformers agentic tool loop.
Wraps a single-turn cumulative-text generator (the existing
``InferenceOrchestrator.generate_chat_response`` pipeline that streams
from a worker subprocess) with the tool-calling, thinking-block,
status, and metadata event protocol used by the GGUF path. Keeps the
front-end SSE shape identical across backends so the chat UI does not
care which engine actually ran the model.
The GGUF path lives in ``llama_cpp.py`` and talks to llama-server's
structured ``delta.tool_calls`` directly. Native transformers has no
such structured channel, so this loop parses tool calls from the
cumulative text and dispatches them via ``core.inference.tools``.
"""
import json
import threading
from typing import Callable, Generator, Optional
from urllib.parse import urlparse
from loggers import get_logger
from core.inference.tool_call_parser import (
BUDGET_EXHAUSTED_NUDGE,
DUPLICATE_CALL_NUDGE,
TOOL_ERROR_NUDGE,
TOOL_ERROR_PREFIXES,
TOOL_XML_SIGNALS,
has_tool_signal,
parse_tool_calls_from_text,
strip_tool_markup,
)
logger = get_logger(__name__)
# Buffer cap while waiting to disambiguate a possible tool-call prefix.
_MAX_BUFFER_CHARS = 32
def _status_for_tool(tool_name: str, arguments: dict) -> str:
"""Return a human-readable status line matching the GGUF path."""
if tool_name == "web_search":
url = (arguments.get("url") or "").strip()
if url:
parsed = urlparse(url)
if parsed.scheme in ("http", "https") and parsed.hostname:
host = parsed.hostname
if host.startswith("www."):
host = host[4:]
return f"Reading: {host}"
return "Reading page..."
query = arguments.get("query", "")
return f"Searching: {query}"
if tool_name == "python":
preview = (arguments.get("code") or "").strip().split("\n")[0][:60]
return f"Running Python: {preview}" if preview else "Running Python..."
if tool_name == "terminal":
preview = (arguments.get("command") or "")[:60]
return f"Running: {preview}" if preview else "Running command..."
return f"Calling: {tool_name}"
_CANONICAL_HEAL_ARG = {"python": "code", "terminal": "command"}
def _coerce_arguments(raw_args, *, heal: bool, tool_name: str = "") -> dict:
"""Normalise tool ``arguments`` to a dict.
Some templates emit a JSON string, others a bare query string. With
``heal=True`` we accept a bare string as ``{<canonical_key>: ...}``
so a Hermes-style call without proper JSON still runs the tool. The
canonical key is picked per tool: ``code`` for python, ``command``
for terminal, ``query`` for everything else (e.g. web_search).
"""
if isinstance(raw_args, dict):
return raw_args
if isinstance(raw_args, str):
try:
parsed = json.loads(raw_args)
if isinstance(parsed, dict):
return parsed
except (json.JSONDecodeError, ValueError):
pass
if heal:
key = _CANONICAL_HEAL_ARG.get(tool_name, "query")
return {key: raw_args}
return {"raw": raw_args}
return {}
def run_safetensors_tool_loop(
*,
single_turn: Callable[[list], Generator[str, None, None]],
messages: list[dict],
tools: list[dict],
execute_tool: Callable[..., str],
cancel_event: Optional[threading.Event] = None,
auto_heal_tool_calls: bool = True,
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
``single_turn(messages)`` must yield cumulative assistant text
(each yield is a snapshot including all previously emitted tokens).
The loop:
* Buffers the leading characters of every turn so it can decide
whether the model is about to emit a tool call. Plain content
starts streaming as soon as the buffer rules it out.
* On detecting ``<tool_call>`` or ``<function=`` in the cumulative
text, drains the rest of the turn silently and parses tool calls
out of the full content.
* Executes each tool via ``execute_tool``, appends the assistant
tool-call message and the tool result to the conversation, and
re-enters ``single_turn`` for the next iteration.
* After ``max_tool_iterations`` turns without a final answer, asks
the model once more to produce a final answer with no tools.
Yields event dicts matching the GGUF path:
* ``{"type": "status", "text": ...}`` -- empty string clears the badge.
* ``{"type": "content", "text": ...}`` -- cumulative cleaned text for
the current assistant turn (the consumer should diff against its
own ``prev_text`` cursor).
* ``{"type": "tool_start", "tool_name", "tool_call_id", "arguments"}``
* ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}``
"""
conversation = list(messages)
tool_call_history: list[tuple[str, bool]] = []
final_attempt_done = False
allowed_tool_names = {
(tool.get("function") or {}).get("name")
for tool in (tools or [])
if (tool.get("function") or {}).get("name")
}
next_call_id = 0
if max_tool_iterations <= 0:
# 0 = disabled (same contract as the GGUF loop).
yield {"type": "status", "text": ""}
return
_state_buffering = 0
_state_streaming = 1
_state_draining = 2
for iteration in range(max_tool_iterations + 1):
if cancel_event is not None and cancel_event.is_set():
return
detect_state = _state_buffering
content_buffer = ""
content_accum = ""
cumulative_display = ""
last_emitted = ""
gen = single_turn(conversation)
prev_cumulative = ""
for cumulative in gen:
if cancel_event is not None and cancel_event.is_set():
return
if not isinstance(cumulative, str):
continue # defensive: pipeline only yields strings
delta = cumulative[len(prev_cumulative) :]
prev_cumulative = cumulative
if not delta:
continue
content_accum += delta
if detect_state == _state_draining:
continue
if detect_state == _state_streaming:
candidate = cumulative_display + delta
signal_pos = -1
for sig in TOOL_XML_SIGNALS:
p = candidate.find(sig)
if p >= 0 and (signal_pos < 0 or p < signal_pos):
signal_pos = p
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
cleaned_before = strip_tool_markup(before_tool)
if len(cleaned_before) > len(last_emitted):
last_emitted = cleaned_before
yield {"type": "content", "text": cleaned_before}
cumulative_display = candidate
detect_state = _state_draining
continue
cumulative_display = candidate
cleaned = strip_tool_markup(cumulative_display)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
continue
# BUFFERING: hold until we know it is not a tool call.
content_buffer += delta
stripped = content_buffer.lstrip()
if not stripped:
continue
is_match = False
is_prefix = False
for sig in TOOL_XML_SIGNALS:
if stripped.startswith(sig):
is_match = True
break
if sig.startswith(stripped):
is_prefix = True
break
if is_match:
detect_state = _state_draining
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
continue
else:
detect_state = _state_streaming
cumulative_display += content_buffer
cleaned = strip_tool_markup(cumulative_display)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Stream finished -- resolve what we collected.
if cancel_event is not None and cancel_event.is_set():
return
if detect_state == _state_buffering:
# Buffer never resolved -- tool XML or plain content.
stripped = content_buffer.lstrip()
if stripped and has_tool_signal(stripped):
detect_state = _state_draining
else:
if content_buffer:
cumulative_display += content_buffer
yield {
"type": "content",
"text": strip_tool_markup(cumulative_display, final = True),
}
yield {"type": "status", "text": ""}
return
if detect_state == _state_streaming:
# No tool detected mid-stream -- check for late tool XML.
safety_tc = None
if has_tool_signal(content_accum):
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
)
if not safety_tc:
# Final answer: streaming already emitted content.
# Skip a final=True re-strip so literal "<tool_call>"
# in prose survives when no real tool call parsed.
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
content_text = strip_tool_markup(content_accum, final = True)
logger.info(
"Safetensors safety net: parsed %d tool call(s) from streamed content",
len(tool_calls),
)
else:
# DRAINING: parse tool calls out of full content.
tool_calls = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
)
if not tool_calls and auto_heal_tool_calls:
# Parser found nothing -- surface raw content so any
# literal "<tool_call>" prose is preserved.
if content_accum:
yield {"type": "content", "text": content_accum}
yield {"type": "status", "text": ""}
return
content_text = strip_tool_markup(content_accum, final = True)
if final_attempt_done:
# Final-answer turn re-called a tool -- stop the loop.
if content_text:
yield {"type": "content", "text": content_text}
yield {"type": "status", "text": ""}
return
assistant_msg: dict = {"role": "assistant", "content": content_text}
if tool_calls:
assistant_msg["tool_calls"] = tool_calls
next_call_id += len(tool_calls)
conversation.append(assistant_msg)
for tc in tool_calls or []:
func = tc.get("function", {}) or {}
tool_name = func.get("name", "") or ""
arguments = _coerce_arguments(
func.get("arguments", {}),
heal = auto_heal_tool_calls,
tool_name = tool_name,
)
yield {"type": "status", "text": _status_for_tool(tool_name, arguments)}
yield {
"type": "tool_start",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"arguments": arguments,
}
tc_key = tool_name + str(arguments)
if allowed_tool_names and tool_name not in allowed_tool_names:
result = (
f"Error: tool '{tool_name}' is not enabled for this "
"request. Use one of the enabled tools or provide a "
"final answer."
)
else:
already_ran_ok = any(
k == tc_key and not err for k, err in tool_call_history
)
if already_ran_ok:
result = DUPLICATE_CALL_NUDGE
else:
eff_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
)
try:
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
yield {
"type": "tool_end",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"result": result,
}
is_error = isinstance(result, str) and result.lstrip().startswith(
TOOL_ERROR_PREFIXES
)
tool_call_history.append((tc_key, is_error))
# Strip frontend image sentinel from the model's view.
# Cut at the first occurrence so leading and consecutive
# sentinels are both removed.
result_for_model = result
if isinstance(result_for_model, str) and "__IMAGES__:" in result_for_model:
result_for_model = result_for_model.split("__IMAGES__:", 1)[0].rstrip()
if is_error:
result_for_model = result_for_model + TOOL_ERROR_NUDGE
tool_msg: dict = {
"role": "tool",
"name": tool_name,
"content": result_for_model,
}
tool_call_id = tc.get("id")
if tool_call_id:
tool_msg["tool_call_id"] = tool_call_id
conversation.append(tool_msg)
# Clear the status badge before the next turn.
yield {"type": "status", "text": ""}
if iteration + 1 >= max_tool_iterations and not final_attempt_done:
# Budget exhausted; nudge a final plain answer.
final_attempt_done = True
conversation.append(
{
"role": "user",
"content": BUDGET_EXHAUSTED_NUDGE,
}
)
yield {"type": "status", "text": ""}

204
studio/backend/core/inference/tool_call_parser.py Normal file
View file

@ -0,0 +1,204 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Backend-neutral tool-call XML parser shared by GGUF and safetensors.
Tolerates missing closing tags in either ``<tool_call>{json}</tool_call>``
or ``<function=name><parameter=k>v...`` shape.
"""
import json
import re
# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing
# unclosed runs so truncated tails don't leak markup.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
]
# Prefixes the streaming buffer watches for to gate in-progress text.
TOOL_XML_SIGNALS = ("<tool_call>", "<function=")
# Nudges + error prefixes shared by the GGUF and safetensors loops.
TOOL_ERROR_PREFIXES = (
"Error",
"Search failed",
"Execution error",
"Blocked:",
"Exit code",
"Failed to fetch",
"Failed to resolve",
"No query provided",
)
DUPLICATE_CALL_NUDGE = (
"You already made this exact call. Do not repeat the same tool "
"call. Try a different approach: fetch a URL from previous "
"results, use Python to process data you already have, or "
"provide your final answer now."
)
TOOL_ERROR_NUDGE = (
"\n\nThe tool call encountered an issue. Please try a different "
"approach or rephrase your request."
)
BUDGET_EXHAUSTED_NUDGE = (
"You have used all available tool calls. Based on everything you "
"have found so far, provide your final answer now. Do not call "
"any more tools."
)
# Pre-compiled patterns reused by ``parse_tool_calls_from_text``.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
def strip_tool_markup(text: str, *, final: bool = False) -> str:
"""Strip tool-call XML from streamed text.
``final=False`` only removes closed pairs (used during streaming so
in-progress XML stays buffered). ``final=True`` also removes a
trailing unclosed run and trims the result.
"""
pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
for pat in pats:
text = pat.sub("", text)
return text.strip() if final else text
def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict]:
"""Parse OpenAI-format ``tool_calls`` from model text.
Returns a list of ``{"id", "type", "function": {"name", "arguments"}}``
dicts. ``arguments`` is always a JSON string so callers can hand it
straight back into an OpenAI-style response.
Handles two shapes:
- JSON inside ``<tool_call>`` tags:
``<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>``
- XML-style function blocks:
``<function=name><parameter=k>v</parameter></function>``
Closing tags (``</tool_call>``, ``</function>``, ``</parameter>``)
are all optional since models frequently omit them.
"""
tool_calls: list[dict] = []
# Pattern 1: <tool_call>{json}. Balanced-brace scan that skips
# braces inside JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
ch = content[i]
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
break
i += 1
if depth == 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": obj.get("name", ""),
"arguments": obj.get("arguments", {}),
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(
tc["function"]["arguments"]
)
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
# Pattern 2: <function=name><parameter=k>v... -- closing tags
# optional; don't use </function> as body boundary because code
# values can contain that literal.
if not tool_calls:
func_starts = list(_TC_FUNC_START_RE.finditer(content))
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
next_func = (
func_starts[idx + 1].start()
if idx + 1 < len(func_starts)
else len(content)
)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()
else:
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
body = _TC_FUNC_CLOSE_RE.sub("", body)
arguments: dict = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
# Single param: take everything to body end so
# embedded </parameter> in code strings is preserved.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
else:
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
next_param = (
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
val = body[val_start:next_param]
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": func_name,
"arguments": json.dumps(arguments),
},
}
tool_calls.append(tc)
return tool_calls
def has_tool_signal(text: str) -> bool:
"""Return True if ``text`` contains any tool-call XML signal."""
return any(s in text for s in TOOL_XML_SIGNALS)

62
studio/backend/core/inference/worker.py
View file

@ -346,6 +346,26 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
# Forward chat_template_info so the parent can classify
# capabilities without re-entering the subprocess.
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
_tpl_info = _entry.get("chat_template_info")
if isinstance(_tpl_info, dict):
model_info["chat_template_info"] = {
"has_template": bool(_tpl_info.get("has_template", False)),
"template": _tpl_info.get("template"),
"format_type": _tpl_info.get("format_type", "generic"),
"template_name": _tpl_info.get("template_name"),
"special_tokens": _tpl_info.get("special_tokens", {}) or {},
}
except Exception as _tpl_exc:
logger.warning("chat_template_info forward failed: %s", _tpl_exc)
_send_response(
resp_queue,
{
@ -416,6 +436,18 @@ def _handle_generate(
"cancel_event": cancel_event,
}
# Optional template/tool plumbing: only forward keys that are
# actually present so the backend signature can evolve without
# breaking older command payloads.
for opt_key in (
"tools",
"enable_thinking",
"reasoning_effort",
"preserve_thinking",
):
if opt_key in cmd:
gen_kwargs[opt_key] = cmd[opt_key]
# Choose generation path
use_adapter = cmd.get("use_adapter")
if use_adapter is not None:
@ -648,36 +680,6 @@ def run_inference_process(
os.environ["HF_HUB_DISABLE_XET"] = "1"
logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)")
# Offline auto-detect: skip 25s of hf_hub_download retries per file
# if DNS is dead; cached files resolve instantly under HF_HUB_OFFLINE=1.
# Scope is this subprocess only -- orchestrator spawns a fresh worker
# per load (see core/inference/orchestrator.py), so the env cannot
# persist across loads.
if "HF_HUB_OFFLINE" not in os.environ:
import socket as _socket
import threading as _threading
# Probe on a daemon thread so concurrent sockets in the parent
# interpreter are not affected by socket.setdefaulttimeout.
_result: list = [None]
def _probe() -> None:
try:
_socket.gethostbyname("huggingface.co")
_result[0] = False
except Exception:
_result[0] = True
_t = _threading.Thread(target = _probe, daemon = True)
_t.start()
_t.join(2.0)
if _result[0] is None or _result[0] is True:
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
logger.warning(
"huggingface.co unreachable; HF_HUB_OFFLINE=1 set for this worker."
)
import warnings
from loggers.config import LogConfig

173
studio/backend/core/tool_healing.py Normal file
View file

@ -0,0 +1,173 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tool-call XML parsing and stripping helpers.
Extracted verbatim from studio/backend/core/inference/llama_cpp.py so that
external inference servers (llama-server wrappers, llama-swap, custom
shims) can reuse the same logic without importing the inference
orchestrator, structlog, httpx, or the rest of the studio backend.
The regexes and function bodies are byte-for-byte identical to the
original inline implementation in llama_cpp.py. Any change made here must
preserve that equivalence; tests/python/test_tool_healing_extraction_is_exact.py
verifies it with AST comparison.
"""
import json
import re
# Pre-compiled patterns for tool XML stripping.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
]
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
def parse_tool_calls_from_text(content: str) -> list[dict]:
"""
Parse tool calls from XML markup in content text.
Handles formats like:
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
Closing tags (</tool_call>, </function>, </parameter>) are all optional
since models frequently omit them.
"""
tool_calls = []
# Pattern 1: JSON inside <tool_call> tags.
# Use balanced-brace extraction that skips braces inside JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
ch = content[i]
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2 # skip escaped character
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
break
i += 1
if depth == 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {
"name": obj.get("name", ""),
"arguments": obj.get("arguments", {}),
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(
tc["function"]["arguments"]
)
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
# Pattern 2: XML-style <function=name><parameter=key>value</parameter></function>
# All closing tags optional -- models frequently omit </parameter>,
# </function>, and/or </tool_call>.
if not tool_calls:
# Step 1: Find all <function=name> positions and extract their bodies.
# Body boundary: use only </tool_call> or next <function= as hard
# boundaries. We avoid using </function> as a boundary because
# code parameter values can contain that literal string.
# After extracting, we trim a trailing </function> if present.
func_starts = list(_TC_FUNC_START_RE.finditer(content))
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
# Hard boundaries: next <function= tag or </tool_call>
next_func = (
func_starts[idx + 1].start()
if idx + 1 < len(func_starts)
else len(content)
)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()
else:
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
# Trim trailing </function> if present (it's the real closing tag)
body = _TC_FUNC_CLOSE_RE.sub("", body)
# Step 2: Extract parameters from body.
# For single-parameter functions (the common case: code, command,
# query), use body end as the only boundary to avoid false matches
# on </parameter> inside code strings.
arguments = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
# Single parameter: value is everything from after the tag
# to end of body, trimming any trailing </parameter>.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
else:
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
# Value ends at next <parameter= or end of body
next_param = (
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
val = body[val_start:next_param]
# Trim trailing </parameter> if present
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
tc = {
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {
"name": func_name,
"arguments": json.dumps(arguments),
},
}
tool_calls.append(tc)
return tool_calls
def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
"""Strip tool-call XML markup from text.
When ``final`` is False, only fully closed tool-call blocks are removed.
When ``final`` is True, trailing incomplete tool-call blocks are removed
too, and the result is stripped of surrounding whitespace.
"""
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
for pat in patterns:
text = pat.sub("", text)
return text.strip() if final else text

51
studio/backend/models/inference.py
View file

@ -70,7 +70,28 @@ class LoadRequest(BaseModel):
)
speculative_type: Optional[str] = Field(
None,
description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.",
description = (
"Speculative decoding mode for GGUF models. Canonical values: "
"'auto' (platform-aware: MTP on MTP GGUFs, ngram-mod fallback "
"for sub-3B), 'mtp' (force draft-mtp only on both GPU and CPU), "
"'ngram' (force ngram-mod only), 'mtp+ngram' (force "
"ngram-mod+draft-mtp chain on both platforms), 'off' (disabled). "
"Legacy values 'default' (-> auto), 'draft-mtp' (-> mtp), "
"'ngram-mod' (-> ngram), and 'ngram-simple' (kept as-is) are "
"still accepted. Ignored for non-GGUF and vision models."
),
)
spec_draft_n_max: Optional[int] = Field(
None,
ge = 1,
le = 16,
description = (
"Max draft tokens per step for MTP speculative decoding "
"(--spec-draft-n-max). Defaults to 2 on GPU and 3 on CPU/Mac "
"when unset (upstream-bench sweet spot for dense Qwen3.6 MTP "
"quants). Only applied when speculative_type resolves to "
"'mtp' or 'mtp+ngram'."
),
)
llama_extra_args: Optional[List[str]] = Field(
None,
@ -218,7 +239,19 @@ class LoadResponse(BaseModel):
)
speculative_type: Optional[str] = Field(
None,
description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
description = (
"Canonical UI-facing requested speculative decoding mode "
"('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / "
"'ngram-simple'), round-tripped from the original LoadRequest "
"via _canonicalize_spec_mode. None when no model is loaded."
),
)
spec_draft_n_max: Optional[int] = Field(
None,
description = (
"Active --spec-draft-n-max for MTP speculative decoding, or "
"None when the platform default is in effect."
),
)
@ -340,7 +373,19 @@ class InferenceStatusResponse(BaseModel):
)
speculative_type: Optional[str] = Field(
None,
description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
description = (
"Canonical UI-facing requested speculative decoding mode "
"('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / "
"'ngram-simple'), round-tripped from the original LoadRequest. "
"None when no model is loaded."
),
)
spec_draft_n_max: Optional[int] = Field(
None,
description = (
"Active --spec-draft-n-max for MTP speculative decoding, or "
"None when the platform default is in effect."
),
)
llama_cpp_supports_mtp: bool = Field(
True,

553
studio/backend/routes/inference.py
View file

@ -117,6 +117,7 @@ try:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_canonicalize_spec_mode,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
@ -143,6 +144,7 @@ except ImportError:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_canonicalize_spec_mode,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
@ -233,6 +235,57 @@ router = APIRouter()
studio_router = APIRouter()
def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
"""Classify reasoning/tool capabilities via the GGUF classifier so
flags match across backends. gpt-oss is overridden because Harmony
routes reasoning and tools through tokenizer channels, not template
markup."""
model_id = getattr(backend, "active_model_name", None)
flags = (
detect_reasoning_flags(
chat_template,
model_identifier = model_id,
log_source = "safetensors",
)
if chat_template
else {
"supports_reasoning": False,
"reasoning_style": "enable_thinking",
"reasoning_always_on": False,
"supports_preserve_thinking": False,
"supports_tools": False,
}
)
# Our safetensors loop only parses <tool_call>{json}</tool_call>
# and <function=name>...</function>. Llama uses <|python_tag|>,
# Mistral uses [TOOL_CALLS]; advertising tools for those would
# enable a pill the parser cannot honour. GGUF is unaffected --
# llama-server normalises every format into structured deltas.
if (
flags.get("supports_tools")
and chat_template
and "<tool_call>" not in chat_template
and "<function=" not in chat_template
):
logger.info(
"safetensors: template advertises tools but uses an "
"emission format the loop cannot parse; suppressing "
"supports_tools"
)
flags["supports_tools"] = False
# gpt-oss: keep reasoning on, drop tools (Harmony channel, not
# <tool_call> XML this loop parses).
try:
if hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model():
flags["supports_reasoning"] = True
flags["reasoning_style"] = "reasoning_effort"
flags["supports_tools"] = False
except Exception:
logger.debug("gpt_oss_check_failed", exc_info = True)
return flags
def _effective_enable_tools(payload) -> Optional[bool]:
"""Resolve `payload.enable_tools` against the process-level tool policy.
@ -441,12 +494,17 @@ def _request_matches_loaded_settings(
# spec on ``not is_vision``), so treat the request as ``off`` against
# the backend's ``None`` to avoid forcing a redundant reload.
if llama_backend.is_vision:
req_spec = "off"
req_mode = "off"
else:
req_spec = _normalise_settings_str(request.speculative_type) or "off"
backend_spec = _normalise_settings_str(llama_backend.speculative_type) or "off"
if req_spec != backend_spec:
req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto"
backend_mode = llama_backend.requested_spec_mode or "auto"
if req_mode != backend_mode:
return False
# spec_draft_n_max only matters when an MTP variant is engaged; None
# means "platform default" and matches whatever the backend chose.
if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None:
if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0):
return False
if (request.chat_template_override or None) != (
llama_backend.chat_template_override or None
):
@ -583,8 +641,10 @@ async def load_model(
reasoning_style = llama_backend.reasoning_style,
reasoning_always_on = llama_backend.reasoning_always_on,
supports_preserve_thinking = llama_backend.supports_preserve_thinking,
supports_tools = llama_backend.supports_tools,
chat_template = llama_backend.chat_template,
speculative_type = llama_backend.speculative_type,
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
)
else:
if (
@ -604,21 +664,10 @@ async def load_model(
logger.warning(
f"Could not retrieve chat template for {backend.active_model_name}: {e}"
)
# Non-GGUF: only advertise reasoning for gpt-oss Harmony,
# which emits reasoning via channels at the tokenizer level.
# Template-level chat_template_kwargs (enable_thinking /
# preserve_thinking / tools) are not yet forwarded through
# the transformers generation path, so avoid advertising
# controls the server cannot honour outside GGUF.
_sf_supports_reasoning = False
_sf_reasoning_style = "enable_thinking"
if hasattr(backend, "_is_gpt_oss_model"):
try:
if backend._is_gpt_oss_model():
_sf_supports_reasoning = True
_sf_reasoning_style = "reasoning_effort"
except Exception:
pass
# Classify via the same path as GGUF.
_sf_flags = _detect_safetensors_features(backend, _chat_template)
_sf_supports_reasoning = _sf_flags["supports_reasoning"]
_sf_reasoning_style = _sf_flags["reasoning_style"]
return LoadResponse(
status = "already_loaded",
model = model_log_label
@ -639,9 +688,9 @@ async def load_model(
),
supports_reasoning = _sf_supports_reasoning,
reasoning_style = _sf_reasoning_style,
reasoning_always_on = False,
supports_preserve_thinking = False,
supports_tools = False,
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
chat_template = _chat_template,
)
@ -724,7 +773,10 @@ async def load_model(
llama_backend.extra_args,
strip_context = "max_seq_length" in fields_set,
strip_cache = "cache_type_kv" in fields_set,
strip_spec = "speculative_type" in fields_set,
strip_spec = (
"speculative_type" in fields_set
or "spec_draft_n_max" in fields_set
),
strip_template = "chat_template_override" in fields_set,
)
try:
@ -765,6 +817,7 @@ async def load_model(
chat_template_override = request.chat_template_override,
cache_type_kv = request.cache_type_kv,
speculative_type = request.speculative_type,
spec_draft_n_max = request.spec_draft_n_max,
n_parallel = _n_parallel,
extra_args = extra_llama_args,
)
@ -788,6 +841,7 @@ async def load_model(
chat_template_override = request.chat_template_override,
cache_type_kv = request.cache_type_kv,
speculative_type = request.speculative_type,
spec_draft_n_max = request.spec_draft_n_max,
n_parallel = _n_parallel,
extra_args = extra_llama_args,
)
@ -846,7 +900,8 @@ async def load_model(
supports_tools = llama_backend.supports_tools,
cache_type_kv = llama_backend.cache_type_kv,
chat_template = llama_backend.chat_template,
speculative_type = llama_backend.speculative_type,
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
)
# ── Standard path: load via Unsloth/transformers ──────────
@ -968,19 +1023,8 @@ async def load_model(
except Exception:
pass
# Non-GGUF: gpt-oss Harmony surfaces reasoning via tokenizer-level
# channels; other safetensors reasoning/tools/preserve-thinking
# knobs are not forwarded to tokenizer.apply_chat_template yet, so
# we only advertise support for the Harmony case here.
_sf_supports_reasoning = False
_sf_reasoning_style = "enable_thinking"
if hasattr(backend, "_is_gpt_oss_model"):
try:
if backend._is_gpt_oss_model():
_sf_supports_reasoning = True
_sf_reasoning_style = "reasoning_effort"
except Exception:
pass
# Classify reasoning/tool flags via the GGUF sniffer.
_sf_flags = _detect_safetensors_features(backend, _chat_template)
return LoadResponse(
status = "loaded",
@ -998,11 +1042,11 @@ async def load_model(
requires_trust_remote_code = bool(
inference_config.get("trust_remote_code", False)
),
supports_reasoning = _sf_supports_reasoning,
reasoning_style = _sf_reasoning_style,
reasoning_always_on = False,
supports_preserve_thinking = False,
supports_tools = False,
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
chat_template = _chat_template,
)
@ -1345,7 +1389,8 @@ async def get_status(
native_context_length = llama_backend.native_context_length,
cache_type_kv = llama_backend.cache_type_kv,
chat_template_override = llama_backend.chat_template_override,
speculative_type = llama_backend.speculative_type,
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
llama_cpp_supports_mtp = _supports_mtp,
llama_cpp_prebuilt_stale = _stale,
llama_cpp_installed_tag = _installed_tag,
@ -1373,18 +1418,8 @@ async def get_status(
else None
)
# Non-GGUF: only gpt-oss Harmony is wired through the transformers
# generation path. Other template-level reasoning / tool kwargs
# are not yet forwarded, so we do not advertise them here.
supports_reasoning = False
reasoning_style = "enable_thinking"
if backend.active_model_name and hasattr(backend, "_is_gpt_oss_model"):
try:
if backend._is_gpt_oss_model():
supports_reasoning = True
reasoning_style = "reasoning_effort"
except Exception:
pass
# Non-GGUF: classify from the loaded template.
_sf_flags = _detect_safetensors_features(backend, chat_template)
inference_config = (
load_inference_config(backend.active_model_name)
if backend.active_model_name
@ -1404,11 +1439,11 @@ async def get_status(
requires_trust_remote_code = bool(
(inference_config or {}).get("trust_remote_code", False)
),
supports_reasoning = supports_reasoning,
reasoning_style = reasoning_style,
reasoning_always_on = False,
supports_preserve_thinking = False,
supports_tools = False,
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
chat_template = chat_template,
llama_cpp_supports_mtp = _supports_mtp,
llama_cpp_prebuilt_stale = _stale,
@ -2254,46 +2289,11 @@ async def openai_chat_completions(
detail = "Audio input is not supported for GGUF chat models yet.",
)
# Reject images if this GGUF model doesn't support vision
image_b64 = extracted_image_b64 or payload.image_base64
if image_b64 and not llama_backend.is_vision:
raise HTTPException(
status_code = 400,
detail = "Image provided but current GGUF model does not support vision.",
)
# Convert image to PNG for llama-server (stb_image has limited format support)
if image_b64:
try:
import base64 as _b64
from io import BytesIO as _BytesIO
from PIL import Image as _Image, UnidentifiedImageError as _UIE
raw = _b64.b64decode(image_b64)
# Normalize to RGB so PNG encoding succeeds regardless of
# source mode (RGBA, P, L, CMYK, I, F, ...). Previously
# we only converted RGBA, which left CMYK/I/F to raise at
# img.save(PNG).
img = _Image.open(_BytesIO(raw)).convert("RGB")
buf = _BytesIO()
img.save(buf, format = "PNG")
image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii")
except _UIE:
raise HTTPException(
status_code = 400,
detail = "Unsupported or corrupt image format.",
)
except Exception:
raise HTTPException(
status_code = 400,
detail = "Failed to process image.",
)
# Build message list with system prompt prepended
gguf_messages = []
if system_prompt:
gguf_messages.append({"role": "system", "content": system_prompt})
gguf_messages.extend(chat_messages)
gguf_messages, has_gguf_image = _openai_messages_for_gguf_chat(
payload,
llama_backend.is_vision,
)
image_b64 = None
cancel_event = threading.Event()
@ -2307,7 +2307,7 @@ async def openai_chat_completions(
use_tools = (
_effective_enable_tools(payload)
and llama_backend.supports_tools
and not image_b64
and not has_gguf_image
)
if use_tools:
@ -2769,6 +2769,300 @@ async def openai_chat_completions(
except Exception as e:
raise HTTPException(status_code = 400, detail = f"Failed to decode image: {e}")
# Classify capability flags from the loaded template.
_sf_model_info = backend.models.get(backend.active_model_name, {})
_sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template")
_sf_features = _detect_safetensors_features(backend, _sf_tpl)
cancel_event = threading.Event()
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
# ── Safetensors tool-calling path ─────────────────────────
# Mirrors the GGUF agentic loop's event shape. Disabled for
# vision turns (untested overlap with image render slot) and
# for gpt-oss (Harmony uses dedicated channels, not <tool_call>
# XML -- gpt-oss tools still work via the GGUF path).
_sf_is_gptoss = False
try:
_sf_is_gptoss = bool(
hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model()
)
except Exception:
_sf_is_gptoss = False
_sf_tool_budget = (
payload.max_tool_calls_per_message
if payload.max_tool_calls_per_message is not None
else 25
)
_sf_use_tools = (
_effective_enable_tools(payload)
and _sf_features.get("supports_tools", False)
and image is None
and not _sf_is_gptoss
and _sf_tool_budget > 0
)
if _sf_use_tools:
from core.inference.tools import ALL_TOOLS
if payload.enabled_tools is not None:
_sf_tools_to_use = [
t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools
]
else:
_sf_tools_to_use = ALL_TOOLS
_sf_tool_names = {t["function"]["name"] for t in _sf_tools_to_use}
_sf_has_web = "web_search" in _sf_tool_names
_sf_has_code = "python" in _sf_tool_names or "terminal" in _sf_tool_names
_sf_date_line = f"The current date is {_date.today().isoformat()}."
_sf_model_size_b = _extract_model_size_b(model_name)
_sf_is_small_model = _sf_model_size_b is not None and _sf_model_size_b < 9
if _sf_is_small_model:
_sf_web_tips = "Do not repeat the same search query."
else:
_sf_web_tips = (
"When you search and find a relevant URL in the results, "
"fetch its full content by calling web_search with the url parameter. "
"Do not repeat the same search query. If a search returns "
"no useful results, try rephrasing or fetching a result URL directly."
)
_sf_code_tips = (
"Use code execution for math, calculations, data processing, "
"or to parse and analyze information from tool results."
)
if _sf_has_web and _sf_has_code:
_sf_nudge = (
_sf_date_line + " "
"You have access to tools. When appropriate, prefer using "
"tools rather than answering from memory. "
+ _sf_web_tips
+ " "
+ _sf_code_tips
)
elif _sf_has_code:
_sf_nudge = (
_sf_date_line + " "
"You have access to tools. When appropriate, prefer using "
"code execution rather than answering from memory. " + _sf_code_tips
)
elif _sf_has_web:
_sf_nudge = (
_sf_date_line + " "
"You have access to tools. When appropriate, prefer using "
"web search for up-to-date or uncertain factual "
"information rather than answering from memory. " + _sf_web_tips
)
else:
_sf_nudge = ""
_sf_system_prompt = system_prompt
if _sf_nudge:
_sf_nudge += _TOOL_ACTION_NUDGE
if _sf_system_prompt:
_sf_system_prompt = _sf_system_prompt.rstrip() + "\n\n" + _sf_nudge
else:
_sf_system_prompt = _sf_nudge
# Strip stale tool-call XML from prior assistant turns.
_sf_chat_messages = []
for _msg in chat_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
_sf_chat_messages.append(
{
**_msg,
"content": _TOOL_XML_RE.sub("", _msg["content"]).strip(),
}
)
else:
_sf_chat_messages.append(_msg)
def sf_generate_with_tools():
return backend.generate_chat_completion_with_tools(
messages = _sf_chat_messages,
tools = _sf_tools_to_use,
system_prompt = _sf_system_prompt or "",
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
max_tokens = payload.max_tokens,
repetition_penalty = payload.repetition_penalty,
cancel_event = cancel_event,
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
preserve_thinking = payload.preserve_thinking,
auto_heal_tool_calls = payload.auto_heal_tool_calls
if payload.auto_heal_tool_calls is not None
else True,
max_tool_iterations = _sf_tool_budget,
tool_call_timeout = payload.tool_call_timeout
if payload.tool_call_timeout is not None
else 300,
session_id = payload.session_id,
use_adapter = payload.use_adapter,
)
_sf_tool_sentinel = object()
_sf_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_sf_tracker = _TrackedCancel(cancel_event, *_sf_cancel_keys)
_sf_tracker.__enter__()
async def sf_tool_stream():
try:
first_chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(role = "assistant"),
finish_reason = None,
)
],
)
yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
gen = sf_generate_with_tools()
prev_text = ""
while True:
if cancel_event.is_set():
backend.reset_generation_state()
break
if await request.is_disconnected():
cancel_event.set()
backend.reset_generation_state()
return
event = await asyncio.to_thread(next, gen, _sf_tool_sentinel)
if event is _sf_tool_sentinel:
break
if event["type"] == "status":
if not event["text"]:
prev_text = ""
status_data = json.dumps(
{
"type": "tool_status",
"content": event["text"],
}
)
yield f"data: {status_data}\n\n"
continue
if event["type"] in ("tool_start", "tool_end"):
if event["type"] == "tool_start":
prev_text = ""
yield f"data: {json.dumps(event)}\n\n"
continue
# Diff cumulative cleaned text against last snapshot.
raw_cumulative = event.get("text", "")
clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
continue
chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(content = new_text),
finish_reason = None,
)
],
)
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
final_chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(),
finish_reason = "stop",
)
],
)
yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
raise
except Exception:
backend.reset_generation_state()
# Generic wire message; full trace stays in the log
# (CWE-209: transformers/torch errors may leak paths).
logger.exception("safetensors tool stream error")
error_chunk = {
"error": {
"message": "An internal error occurred.",
"type": "server_error",
},
}
yield f"data: {json.dumps(error_chunk)}\n\n"
finally:
_sf_tracker.__exit__(None, None, None)
if payload.stream:
return StreamingResponse(
sf_tool_stream(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
# Non-streaming JSON: drain the loop, build one ChatCompletion.
try:
def _drain_to_text():
full_text = ""
gen = sf_generate_with_tools()
for event in gen:
if cancel_event.is_set():
break
if event.get("type") == "content":
full_text = _TOOL_XML_RE.sub("", event.get("text", ""))
return full_text
content_text = await asyncio.to_thread(_drain_to_text)
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = [
CompletionChoice(
message = CompletionMessage(content = content_text),
finish_reason = "stop",
)
],
)
return JSONResponse(content = response.model_dump())
except Exception:
backend.reset_generation_state()
# CWE-209: generic detail; full trace in log.
logger.exception("safetensors tool completion error")
raise HTTPException(
status_code = 500,
detail = "An internal error occurred.",
)
finally:
_sf_tracker.__exit__(None, None, None)
# Shared generation kwargs
gen_kwargs = dict(
messages = chat_messages,
@ -2781,9 +3075,14 @@ async def openai_chat_completions(
max_new_tokens = payload.max_tokens or 2048,
repetition_penalty = payload.repetition_penalty,
)
# Choose generation path (adapter-controlled or standard)
cancel_event = threading.Event()
# Forward reasoning kwargs; the worker/template wrapper peels off
# any the template doesn't accept.
if payload.enable_thinking is not None:
gen_kwargs["enable_thinking"] = payload.enable_thinking
if payload.reasoning_effort is not None:
gen_kwargs["reasoning_effort"] = payload.reasoning_effort
if payload.preserve_thinking is not None:
gen_kwargs["preserve_thinking"] = payload.preserve_thinking
if payload.use_adapter is not None:
@ -2800,9 +3099,6 @@ async def openai_chat_completions(
cancel_event = cancel_event, **gen_kwargs
)
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
# ── Streaming response ────────────────────────────────────────
if payload.stream:
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
@ -4804,6 +5100,47 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
return messages
def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]:
"""Build llama-server messages for the standard GGUF chat path.
llama-server accepts OpenAI multimodal content parts directly. Preserve
all per-turn ``image_url`` parts so multi-image chat history keeps each
image attached to its original turn.
"""
messages = _drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
)
has_message_image = any(
isinstance(msg.get("content"), list)
and any(part.get("type") == "image_url" for part in msg["content"])
for msg in messages
)
if payload.image_base64 and not has_message_image:
# Legacy bytes can be any format; the normalizer below sniffs and
# re-encodes to PNG, so the declared mime is rewritten anyway.
image_part = {
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{payload.image_base64}",
},
}
for msg in reversed(messages):
if msg.get("role") != "user":
continue
existing = msg.get("content")
if isinstance(existing, str):
msg["content"] = [{"type": "text", "text": existing}, image_part]
elif isinstance(existing, list):
existing.append(image_part)
else:
msg["content"] = [image_part]
break
else:
messages.append({"role": "user", "content": [image_part]})
has_image = _normalize_anthropic_openai_images(messages, is_vision)
return messages, has_image
def _extract_response_format(payload):
"""Return the ``response_format`` field on an incoming ChatCompletionRequest
(or None). The model is declared with ``extra="allow"`` so pydantic stashes

1
studio/backend/tests/test_gguf_reload_inheritance.py
View file

@ -78,6 +78,7 @@ def _loaded_backend(**overrides):
backend._requested_n_ctx = 8192
backend._cache_type_kv = None
backend._speculative_type = None
backend._requested_spec_mode = "auto"
backend._chat_template_override = None
backend._is_vision = False
backend._extra_args = None

27
studio/backend/tests/test_kv_cache_estimation.py
View file

@ -1558,6 +1558,33 @@ class TestServerFlags:
)
assert fitted < 32_768
def test_fit_mtp_engaged_returns_smaller_or_equal_context(self):
# MTP-engaged budget is 0.85 of available; non-MTP is 0.90.
# On a tight budget the MTP path must yield <= the non-MTP path.
b = self._gqa_backend()
common = dict(
requested_ctx = 32_768,
available_mib = 128,
model_size_bytes = 8 * 1024 * 1024,
cache_type_kv = "f16",
)
baseline = b._fit_context_to_vram(**common)
mtp = b._fit_context_to_vram(**common, mtp_engaged = True)
assert mtp <= baseline
def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self):
# kv_on_gpu=False short-circuits the fit; mtp_engaged is irrelevant.
b = self._gqa_backend()
fitted = b._fit_context_to_vram(
requested_ctx = 32_768,
available_mib = 1,
model_size_bytes = 100,
cache_type_kv = "f16",
kv_on_gpu = False,
mtp_engaged = True,
)
assert fitted == 32_768
def test_fit_threads_swa_full_through_estimator(self):
# SWA model, generous budget; both should fit but cache size differs.
b = self._swa_backend()

661
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -52,6 +52,9 @@ import pytest
from core.inference.llama_cpp import (
LlamaCppBackend,
_backfill_usage_from_timings,
_build_ngram_mod_flags,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_is_mtp_model_name,
)
@ -186,6 +189,10 @@ def _mtp_backend(**overrides):
backend._requested_n_ctx = 8192
backend._cache_type_kv = None
backend._speculative_type = "draft-mtp"
# Default fixture simulates Auto having auto-promoted to draft-mtp.
# Individual tests override _requested_spec_mode when they want a
# forced mode or the user---spec-type-extra-args path.
backend._requested_spec_mode = "auto"
backend._chat_template_override = None
backend._is_vision = False
backend._extra_args = None
@ -233,9 +240,16 @@ def test_already_in_target_state_matches_when_request_uses_default_for_mtp_model
)
def test_already_in_target_state_non_mtp_model_unaffected():
# Promotion is gated on the name; non-MTP must still mismatch req=None.
backend = _mtp_backend(_model_identifier = "unsloth/Qwen3.6-27B-GGUF")
def test_already_in_target_state_auto_request_matches_auto_backend_for_non_mtp_model():
# Under the requested-mode round-trip model, Auto requested against an
# Auto-recorded backend matches regardless of model name. The underlying
# resolved emission (--spec-default vs draft-mtp) is handled by the
# backend's own load path and reflected in _speculative_type; the
# short-circuit comparison only cares whether the *intent* changed.
backend = _mtp_backend(
_model_identifier = "unsloth/Qwen3.6-27B-GGUF",
_speculative_type = "default",
)
assert (
backend._already_in_target_state(
gguf_path = None,
@ -248,7 +262,7 @@ def test_already_in_target_state_non_mtp_model_unaffected():
extra_args = None,
is_vision = False,
)
is False
is True
)
@ -308,6 +322,7 @@ def test_already_in_target_state_user_spec_type_override_matches_clean_backend()
# User --spec-type none suppressed auto-MTP; repeat /load must not re-promote.
backend = _mtp_backend(
_speculative_type = None,
_requested_spec_mode = None,
_extra_args = ["--spec-type", "none"],
)
assert (
@ -389,12 +404,17 @@ def test_already_in_target_state_vision_mtp_default_matches():
)
def test_already_in_target_state_vision_non_mtp_unaffected():
# Vision non-MTP repo (no -MTP marker) must still mismatch req=None
# against a backend running draft-mtp.
def test_already_in_target_state_vision_off_matches_vision_backend():
# Vision loads silently drop speculative decoding at the route level
# (_request_matches_loaded_settings overrides req to "off"). At the
# llama_cpp.py level, _already_in_target_state compares canonical
# requested modes; a vision backend recorded with _requested_spec_mode
# = "off" matches a req of "off" or None+vision.
backend = _mtp_backend(
_model_identifier = "unsloth/Qwen3-VL-4B-Instruct-GGUF",
_is_vision = True,
_speculative_type = None,
_requested_spec_mode = "off",
)
assert (
backend._already_in_target_state(
@ -403,12 +423,12 @@ def test_already_in_target_state_vision_non_mtp_unaffected():
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
speculative_type = "off",
chat_template_override = None,
extra_args = None,
is_vision = True,
)
is False
is True
)
@ -482,10 +502,17 @@ def _make_fake_llama_server(path: Path, help_text: str) -> Path:
return path
_NEEDS_BASH = pytest.mark.skipif(
sys.platform == "win32",
reason = "fake llama-server is a bash stub; Windows has no direct executor",
)
def _clear_caps_cache():
LlamaCppBackend._capability_cache.clear()
@_NEEDS_BASH
def test_probe_server_capabilities_detects_draft_mtp(tmp_path):
# Original naming from llama.cpp #22673.
fake = _make_fake_llama_server(
@ -500,6 +527,7 @@ def test_probe_server_capabilities_detects_draft_mtp(tmp_path):
assert caps["supports_mtp"] is True
@_NEEDS_BASH
def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
# Renamed upstream: draft-mtp -> mtp.
fake = _make_fake_llama_server(
@ -513,6 +541,7 @@ def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
assert caps["supports_mtp"] is True
@_NEEDS_BASH
def test_probe_server_capabilities_reports_outdated_binary(tmp_path):
# Pre-MTP llama.cpp: only ngram variants.
fake = _make_fake_llama_server(
@ -533,6 +562,130 @@ def test_probe_server_capabilities_handles_missing_binary():
assert caps["supports_mtp"] is False
# ngram-mod flag flavor detection (new vs legacy llama-server).
# Help-text fixtures mirror the actual `llama-server --help` block
# layout (flag on its own line; description indented underneath).
_POST_RENAME_HELP = """\
--spec-draft-n-max N number of tokens to draft for speculative decoding (default: 16)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX)
--spec-draft-n-min N minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN)
--spec-draft-p-min, --draft-p-min P minimum speculative decoding probability (greedy) (default: 0.75)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN)
--spec-ngram-mod-n-min N minimum number of ngram tokens (default: 48)
--spec-ngram-mod-n-max N maximum number of ngram tokens (default: 64)
--spec-ngram-mod-n-match N ngram-mod lookup length (default: 24)
--spec-type none,draft-simple,draft-mtp,ngram-mod comma-separated list of types of speculative decoding to use
(env: LLAMA_ARG_SPEC_TYPE)
--draft, --draft-n, --draft-max N the argument has been removed. use --spec-draft-n-max or --spec-ngram-mod-n-max
(env: LLAMA_ARG_DRAFT_MAX)
--draft-min, --draft-n-min N the argument has been removed. use --spec-draft-n-min or --spec-ngram-mod-n-min
(env: LLAMA_ARG_DRAFT_MIN)
--spec-ngram-size-n N the argument has been removed. use the respective --spec-ngram-*-size-n or --spec-ngram-mod-n-match
"""
_LEGACY_HELP = """\
--draft, --draft-n, --draft-max N number of tokens to draft for speculative decoding (default: 8)
(env: LLAMA_ARG_DRAFT_MAX)
--draft-min, --draft-n-min N minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_DRAFT_MIN)
--spec-ngram-size-n N ngram lookup length (default: 24)
--spec-type none,ngram-mod,ngram-simple comma-separated list of types of speculative decoding to use
"""
@_NEEDS_BASH
def test_probe_detects_post_rename_ngram_mod_flavor(tmp_path):
fake = _make_fake_llama_server(tmp_path / "llama-server", _POST_RENAME_HELP)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["found"] is True
assert caps["ngram_mod_flavor"] == "new"
assert caps["supports_ngram_mod"] is True
assert caps["spec_draft_n_max_flag"] == "--spec-draft-n-max"
@_NEEDS_BASH
def test_probe_detects_legacy_ngram_mod_flavor(tmp_path):
fake = _make_fake_llama_server(tmp_path / "llama-server", _LEGACY_HELP)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["found"] is True
assert caps["ngram_mod_flavor"] == "legacy"
assert caps["supports_ngram_mod"] is True
assert caps["spec_draft_n_max_flag"] == "--draft-max"
@_NEEDS_BASH
def test_probe_ignores_removal_stub_descriptions(tmp_path):
# Post-rename binary: legacy flags are present but with
# "argument has been removed" descriptions; must not be detected
# as legacy.
fake = _make_fake_llama_server(tmp_path / "llama-server", _POST_RENAME_HELP)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["ngram_mod_flavor"] == "new"
@_NEEDS_BASH
def test_probe_no_ngram_mod_on_minimal_binary(tmp_path):
# Pre-anything: neither set present.
fake = _make_fake_llama_server(
tmp_path / "llama-server",
"--spec-type none\n--threads N\n",
)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["ngram_mod_flavor"] is None
assert caps["supports_ngram_mod"] is False
def test_build_ngram_mod_flags_new():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"})
assert flags == [
"--spec-ngram-mod-n-match",
"24",
"--spec-ngram-mod-n-min",
"48",
"--spec-ngram-mod-n-max",
"64",
]
def test_build_ngram_mod_flags_legacy():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "legacy"})
assert flags == [
"--spec-ngram-size-n",
"24",
"--draft-min",
"48",
"--draft-max",
"64",
]
def test_build_ngram_mod_flags_empty_when_unsupported():
assert _build_ngram_mod_flags({"ngram_mod_flavor": None}) == []
assert _build_ngram_mod_flags(None) == []
assert _build_ngram_mod_flags({}) == []
def test_build_ngram_mod_flags_respects_custom_values():
flags = _build_ngram_mod_flags(
{"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32
)
assert flags == [
"--spec-ngram-mod-n-match",
"16",
"--spec-ngram-mod-n-min",
"24",
"--spec-ngram-mod-n-max",
"32",
]
@_NEEDS_BASH
def test_probe_server_capabilities_caches_by_mtime(tmp_path):
# Same (path, mtime) -> cache hit. Bumped mtime -> re-probe.
fake = _make_fake_llama_server(
@ -555,3 +708,493 @@ def test_probe_server_capabilities_caches_by_mtime(tmp_path):
caps2 = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps2["mtp_token"] == "draft-mtp"
assert caps2["supports_mtp"] is True
# spec_draft_n_max plumbing (first-class --spec-draft-n-max override).
def test_already_in_target_state_matches_when_draft_n_max_unset():
# None on the request means "platform default"; matches any backend.
backend = _mtp_backend(_spec_draft_n_max = None)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
spec_draft_n_max = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_matches_when_draft_n_max_equals_backend():
backend = _mtp_backend(_spec_draft_n_max = 4)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
spec_draft_n_max = 4,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_mismatches_when_draft_n_max_differs():
backend = _mtp_backend(_spec_draft_n_max = 4)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
spec_draft_n_max = 8,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is False
)
def test_already_in_target_state_draft_n_max_ignored_when_not_mtp():
# ngram-mod backend; spec_draft_n_max is MTP-only and must not force
# a reload against a non-MTP active spec.
backend = _mtp_backend(
_speculative_type = "ngram-mod",
_requested_spec_mode = "ngram",
_spec_draft_n_max = None,
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "ngram-mod",
spec_draft_n_max = 8,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
# Sub-3B MTP gate -- tiny dense models regress with the MTP draft
# head, so load_model falls back to ngram-mod (when the binary supports
# it) instead of draft-mtp. The reload-skip mirror must follow the
# same fallback so a sub-3B reload-with-default does not bounce a
# correctly-configured ngram-mod / off backend.
def _patch_probe(monkeypatch, ngram_supported):
"""Force probe_server_capabilities to a deterministic result so
tests don't depend on whatever llama-server happens to be on PATH."""
fake = {
"found": True,
"mtp_token": "draft-mtp",
"supports_mtp": True,
"ngram_mod_flavor": "new" if ngram_supported else None,
"supports_ngram_mod": bool(ngram_supported),
"spec_draft_n_max_flag": "--spec-draft-n-max",
}
monkeypatch.setattr(
LlamaCppBackend,
"probe_server_capabilities",
classmethod(lambda cls, binary = None: fake),
)
monkeypatch.setattr(
LlamaCppBackend,
"_find_llama_server_binary",
classmethod(lambda cls: "/fake/llama-server"),
)
def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported(
monkeypatch,
):
# 0.8B MTP request -- load_model would have promoted to ngram-mod
# (no MTP head); reload check must match a ngram-mod backend.
_patch_probe(monkeypatch, ngram_supported = True)
backend = _mtp_backend(
_model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF",
_speculative_type = "ngram-mod",
_spec_draft_n_max = None,
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_sub_3b_falls_back_to_off_when_no_ngram(monkeypatch):
# 0.8B + binary lacks ngram-mod -> fall back to off.
_patch_probe(monkeypatch, ngram_supported = False)
backend = _mtp_backend(
_model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF",
_speculative_type = None,
_spec_draft_n_max = None,
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_4b_mtp_request_promotes_as_before(monkeypatch):
# 4B is above the 3B threshold -> auto-promote still applies.
_patch_probe(monkeypatch, ngram_supported = True)
backend = _mtp_backend(
_model_identifier = "unsloth/Qwen3.5-4B-MTP-GGUF",
_speculative_type = "draft-mtp",
_spec_draft_n_max = None,
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.5-4B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_already_in_target_state_2b_falls_back_to_ngram_below_threshold(monkeypatch):
# 2.0B is below the 3B threshold -> ngram-mod fallback, not
# draft-mtp. Clean-bench shows 2B regresses with draft-mtp.
_patch_probe(monkeypatch, ngram_supported = True)
backend = _mtp_backend(
_model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF",
_speculative_type = "ngram-mod",
_spec_draft_n_max = None,
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
# usage backfill from timings (Studio UI t/s widget fix).
def test_backfill_usage_from_timings_fills_when_completion_tokens_zero():
out = _backfill_usage_from_timings(
{"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
{"prompt_n": 42, "predicted_n": 128, "predicted_per_second": 100.0},
)
assert out["completion_tokens"] == 128
assert out["prompt_tokens"] == 42
assert out["total_tokens"] == 170
def test_backfill_usage_from_timings_fills_when_usage_missing():
out = _backfill_usage_from_timings(
None,
{"prompt_n": 42, "predicted_n": 128, "predicted_per_second": 100.0},
)
assert out["completion_tokens"] == 128
assert out["prompt_tokens"] == 42
assert out["total_tokens"] == 170
def test_backfill_usage_from_timings_preserves_real_usage():
# Non-zero completion_tokens means llama-server reported correctly;
# do not overwrite.
real = {"prompt_tokens": 50, "completion_tokens": 200, "total_tokens": 250}
out = _backfill_usage_from_timings(real, {"predicted_n": 999, "prompt_n": 999})
assert out is real
assert out["completion_tokens"] == 200
def test_backfill_usage_from_timings_passthrough_when_timings_empty():
assert _backfill_usage_from_timings(None, None) is None
assert _backfill_usage_from_timings(None, {}) is None
usage = {"completion_tokens": 0}
# No timings.predicted_n -> nothing to fill, return as-is.
assert _backfill_usage_from_timings(usage, {"prompt_ms": 5.0}) is usage
# ── _canonicalize_spec_mode (pure) ─────────────────────────────────
@pytest.mark.parametrize(
"value, expected",
[
# New canonical values pass through unchanged.
("auto", "auto"),
("mtp", "mtp"),
("ngram", "ngram"),
("mtp+ngram", "mtp+ngram"),
("off", "off"),
("ngram-simple", "ngram-simple"),
# Legacy wire values map onto the new vocabulary.
("default", "auto"),
("draft-mtp", "mtp"),
("ngram-mod", "ngram"),
# Comma-chained legacy values (e.g. from persisted state) collapse
# to the right canonical mode.
("ngram-mod,draft-mtp", "mtp+ngram"),
("draft-mtp,ngram-mod", "mtp+ngram"),
("draft-mtp,mtp", "mtp"),
("ngram-mod,ngram", "ngram"),
# Case and whitespace are ignored.
(" AUTO ", "auto"),
("MTP", "mtp"),
("MTP+Ngram", "mtp+ngram"),
# None / empty / whitespace pass through as None.
(None, None),
("", None),
(" ", None),
# Non-string inputs collapse to None.
(42, None),
(True, None),
# Unknown strings fall back to "auto" (safe default).
("bogus", "auto"),
],
)
def test_canonicalize_spec_mode(value, expected):
assert _canonicalize_spec_mode(value) == expected
# ── _build_speculative_flags resolver matrix ──────────────────────
def _resolver_backend(monkeypatch, *, ngram_supported = True, mtp_token = "draft-mtp"):
"""Backend with a deterministic probe so the resolver is hermetic."""
fake = {
"found": True,
"mtp_token": mtp_token,
"supports_mtp": bool(mtp_token),
"ngram_mod_flavor": "new" if ngram_supported else None,
"supports_ngram_mod": bool(ngram_supported),
"spec_draft_n_max_flag": "--spec-draft-n-max",
}
monkeypatch.setattr(
LlamaCppBackend,
"probe_server_capabilities",
classmethod(lambda cls, binary = None: fake),
)
backend = LlamaCppBackend()
backend._nextn_predict_layers = None
return backend
def _flags_dict(flags):
"""Parse the spec-flag list into a small {flag: value} dict; collapses
repeated flags by keeping the last (only --spec-type can repeat and
never does in our resolver)."""
out = {}
i = 0
while i < len(flags):
token = flags[i]
if i + 1 < len(flags) and not flags[i + 1].startswith("--"):
out[token] = flags[i + 1]
i += 2
else:
out[token] = True
i += 1
return out
_MTP_MODEL = "unsloth/Qwen3.6-27B-MTP-GGUF"
_NON_MTP_MODEL = "unsloth/Qwen3-7B-Instruct-GGUF"
_SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF"
@pytest.mark.parametrize(
"requested, gpus, model, expect_spec_type, expect_n_max, expect_ngram_knobs",
[
# ── auto + MTP model + 3B+: GPU = mtp only, CPU = chain ──
("auto", True, _MTP_MODEL, "draft-mtp", "2", False),
("auto", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True),
# ── auto + non-MTP: emit --spec-default ──
("auto", True, _NON_MTP_MODEL, None, None, False),
("auto", False, _NON_MTP_MODEL, None, None, False),
# ── auto + sub-3B MTP: fallback to ngram-mod ──
("auto", True, _SUB_3B_MTP_MODEL, "ngram-mod", None, True),
("auto", False, _SUB_3B_MTP_MODEL, "ngram-mod", None, True),
# ── mtp forced: MTP-only on BOTH platforms ──
("mtp", True, _MTP_MODEL, "draft-mtp", "2", False),
("mtp", False, _MTP_MODEL, "draft-mtp", "3", False),
# ── mtp forced on sub-3B: engage anyway ──
("mtp", True, _SUB_3B_MTP_MODEL, "draft-mtp", "2", False),
# ── mtp forced on non-MTP: engage anyway ──
("mtp", True, _NON_MTP_MODEL, "draft-mtp", "2", False),
# ── ngram forced: ngram-mod alone on BOTH platforms ──
("ngram", True, _MTP_MODEL, "ngram-mod", None, True),
("ngram", False, _MTP_MODEL, "ngram-mod", None, True),
("ngram", True, _NON_MTP_MODEL, "ngram-mod", None, True),
# ── mtp+ngram forced: chain on BOTH platforms ──
("mtp+ngram", True, _MTP_MODEL, "ngram-mod,draft-mtp", "2", True),
("mtp+ngram", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True),
("mtp+ngram", True, _SUB_3B_MTP_MODEL, "ngram-mod,draft-mtp", "2", True),
# ── off: nothing emitted ──
("off", True, _MTP_MODEL, None, None, False),
("off", False, _MTP_MODEL, None, None, False),
# ── legacy values round-trip to the canonical emission ──
("default", True, _MTP_MODEL, "draft-mtp", "2", False),
("draft-mtp", True, _MTP_MODEL, "draft-mtp", "2", False),
("ngram-mod", True, _MTP_MODEL, "ngram-mod", None, True),
("ngram-mod,draft-mtp", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True),
# ── ngram-simple: pass through ──
("ngram-simple", True, _MTP_MODEL, "ngram-simple", None, False),
],
)
def test_build_speculative_flags_matrix(
monkeypatch,
requested,
gpus,
model,
expect_spec_type,
expect_n_max,
expect_ngram_knobs,
):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = requested,
spec_draft_n_max = None,
extra_args = None,
model_identifier = model,
model_path = None,
gpus = gpus,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
if expect_spec_type is None:
assert "--spec-type" not in parsed
else:
assert parsed.get("--spec-type") == expect_spec_type
if expect_n_max is None:
assert "--spec-draft-n-max" not in parsed
else:
assert parsed.get("--spec-draft-n-max") == expect_n_max
if expect_ngram_knobs:
assert "--spec-ngram-mod-n-match" in parsed
assert "--spec-ngram-mod-n-min" in parsed
assert "--spec-ngram-mod-n-max" in parsed
else:
assert "--spec-ngram-mod-n-match" not in parsed
def test_build_speculative_flags_user_extra_args_owns_spec_type(monkeypatch):
# User --spec-type in extra_args bypasses the dropdown entirely.
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "mtp", # would normally force MTP
spec_draft_n_max = None,
extra_args = ["--spec-type", "ngram-mod"],
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
# No flags emitted by the resolver -- the user's extra_args carries
# the --spec-type, and the resolver records requested_spec_mode = None.
assert flags == []
assert backend.requested_spec_mode is None
assert backend.speculative_type is None
@pytest.mark.parametrize("mode", ["auto", "mtp", "ngram", "mtp+ngram", "off"])
def test_build_speculative_flags_round_trips_requested_mode(monkeypatch, mode):
# The status round-trip is the contract that lets the UI dropdown
# restore its picked value after reload / refresh.
backend = _resolver_backend(monkeypatch)
backend._build_speculative_flags(
speculative_type = mode,
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert backend.requested_spec_mode == mode
def test_build_speculative_flags_user_draft_n_max_override(monkeypatch):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "mtp",
spec_draft_n_max = 5,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-draft-n-max") == "5"
assert backend.spec_draft_n_max == 5
def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
# Outdated llama-server with no MTP support: forced MTP must degrade
# to spec-off (warned) rather than emit a bad --spec-type.
backend = _resolver_backend(monkeypatch, mtp_token = None)
flags = backend._build_speculative_flags(
speculative_type = "mtp",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert "--spec-type" not in flags
# _speculative_type stays None (resolved emission was none), but
# _requested_spec_mode still reflects the user's choice.
assert backend.requested_spec_mode == "mtp"
assert backend.speculative_type is None

6
studio/backend/tests/test_llama_server_args.py
View file

@ -47,17 +47,15 @@ from core.inference.llama_server_args import (
["--spec-type", "draft-mtp", "--spec-draft-n-max", "6"],
[
"--spec-type",
"draft-mtp",
"ngram-mod,draft-mtp",
"--spec-draft-n-max",
"3",
"--spec-type",
"ngram-mod",
"--spec-ngram-mod-n-match",
"24",
"--spec-ngram-mod-n-min",
"48",
"--spec-ngram-mod-n-max",
"6",
"64",
],
# Reasoning controls
["--reasoning-format", "deepseek"],

94
studio/backend/tests/test_mlx_inference_backend.py
View file

@ -158,3 +158,97 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri
assert backend._is_vlm is True
assert isinstance(backend._processor, _DummyProcessor)
assert isinstance(backend._tokenizer, _DummyTokenizer)
# Regression: MLXInferenceBackend.generate_chat_response must accept the
# four template kwargs (tools / enable_thinking / reasoning_effort /
# preserve_thinking) so the route layer can forward what the user
# toggled in the UI. The previous signature raised
# "got an unexpected keyword argument 'tools'" on Mac.
def test_mlx_generate_chat_response_accepts_template_kwargs():
import inspect
from core.inference.mlx_inference import MLXInferenceBackend
sig = inspect.signature(MLXInferenceBackend.generate_chat_response)
params = sig.parameters
for name in ("tools", "enable_thinking", "reasoning_effort", "preserve_thinking"):
assert name in params, (
f"MLX.generate_chat_response is missing the {name!r} kwarg; "
"the route layer forwards this and a missing kwarg raises "
"TypeError on Mac"
)
assert (
params[name].default is None
), f"{name!r} must default to None so existing callers stay valid"
def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
"""The Mac text path must route through apply_chat_template_for_
generation so reasoning / tool kwargs reach the tokenizer."""
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import MLXInferenceBackend
captured = {}
def _fake_apply(tokenizer, messages, **kwargs):
captured["tokenizer"] = tokenizer
captured["messages"] = messages
captured["kwargs"] = kwargs
return "<rendered prompt>"
monkeypatch.setattr(
"core.inference.chat_template_helpers." "apply_chat_template_for_generation",
_fake_apply,
raising = True,
)
# mlx_lm.stream_generate yields response objects with .token; make a
# one-token generator so _generate_text returns without touching the
# real stack.
import types as _types
mlx_lm_pkg = _types.ModuleType("mlx_lm")
mlx_lm_sample = _types.ModuleType("mlx_lm.sample_utils")
mlx_lm_sample.make_sampler = lambda **_kw: object()
mlx_lm_sample.make_logits_processors = lambda **_kw: None
class _Resp:
def __init__(self, tok):
self.token = tok
def _stream_generate(_model, _tokenizer, **_kw):
yield _Resp(1)
mlx_lm_pkg.stream_generate = _stream_generate
monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg)
monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample)
class _Tok:
chat_template = "x"
def decode(self, ids, skip_special_tokens = False):
return "hi"
backend = MLXInferenceBackend()
backend._model = object()
backend._tokenizer = _Tok()
backend._is_vlm = False
out = list(
backend.generate_chat_response(
messages = [{"role": "user", "content": "ping"}],
tools = [{"function": {"name": "web_search"}}],
enable_thinking = True,
reasoning_effort = "medium",
preserve_thinking = True,
max_new_tokens = 1,
)
)
assert out == ["hi"]
# The kwargs the user toggled must reach the chat-template helper.
assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}]
assert captured["kwargs"]["enable_thinking"] is True
assert captured["kwargs"]["reasoning_effort"] == "medium"
assert captured["kwargs"]["preserve_thinking"] is True

109
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -26,6 +26,7 @@ sys.path.insert(0, _backend)
import httpx
import pytest
from fastapi import HTTPException
from pydantic import ValidationError
from models.inference import (
@ -532,6 +533,7 @@ class TestFriendlyErrorHttpx:
from routes.inference import ( # noqa: E402
_drop_empty_assistant_sentinels,
_openai_messages_for_gguf_chat,
_openai_messages_for_passthrough,
)
@ -616,3 +618,110 @@ class TestDropEmptyAssistantSentinels:
assert roles == ["user", "user"]
for m in out:
assert m.get("content"), m
class TestGgufVisionMessages:
_PNG_B64 = (
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADUlEQVR42mNk"
"+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
def test_preserves_multiturn_image_parts_on_original_turns(self):
req = ChatCompletionRequest(
model = "default",
image_base64 = self._PNG_B64,
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "describe image one"},
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{self._PNG_B64}",
},
},
],
},
{"role": "assistant", "content": "first answer"},
{
"role": "user",
"content": [
{"type": "text", "text": "describe image two"},
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{self._PNG_B64}",
},
},
],
},
],
)
messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True)
assert has_image is True
assert messages[0]["content"][0] == {
"type": "text",
"text": "describe image one",
}
assert messages[0]["content"][1]["type"] == "image_url"
assert len(messages[0]["content"]) == 2
assert messages[2]["content"][0] == {
"type": "text",
"text": "describe image two",
}
assert messages[2]["content"][1]["type"] == "image_url"
assert len(messages[2]["content"]) == 2
assert isinstance(messages[1]["content"], str)
# Legacy top-level image_base64 must be ignored when any message-level
# image already exists; otherwise turn 2 ends up with two image parts.
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
image_parts = [p for p in content if p.get("type") == "image_url"]
assert len(image_parts) == 1, msg
def test_legacy_image_base64_is_injected_when_messages_are_text_only(self):
req = ChatCompletionRequest(
model = "default",
image_base64 = self._PNG_B64,
messages = [{"role": "user", "content": "describe this image"}],
)
messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True)
assert has_image is True
assert messages[0]["content"][0] == {
"type": "text",
"text": "describe this image",
}
assert messages[0]["content"][1]["type"] == "image_url"
assert messages[0]["content"][1]["image_url"]["url"].startswith(
"data:image/png;base64,"
)
def test_rejects_image_parts_for_text_only_gguf(self):
req = ChatCompletionRequest(
model = "default",
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "look"},
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{self._PNG_B64}",
},
},
],
},
],
)
with pytest.raises(HTTPException) as exc_info:
_openai_messages_for_gguf_chat(req, is_vision = False)
assert "does not support vision" in str(exc_info.value)

451
studio/backend/tests/test_safetensors_capability_advertise.py Normal file
View file

@ -0,0 +1,451 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Capability advertisement contract: classifier honesty, worker→
orchestrator IPC hop, and route-layer end-to-end. Pure helpers + fakes;
no torch / transformers import.
"""
from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
_backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
# Qwen3 snippet covering tools, enable_thinking, preserve_thinking.
QWEN3_TEMPLATE = """
{%- if tools %}
{{- '<|im_start|>system\\nFor each function call, return a json object'
' wrapped inside <tool_call></tool_call> tags.\\n' }}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
{%- for message in messages %}
{%- if message.role == 'tool' %}
{{- '<|im_start|>tool\\n' + message.content + '<|im_end|>\\n' }}
{%- endif %}
{%- endfor %}
{%- if enable_thinking is defined and enable_thinking %}
{{- '<think>' }}
{%- endif %}
{%- if preserve_thinking %}
{{- assistant.reasoning_content }}
{%- endif %}
"""
GPT_OSS_TEMPLATE = """
<|start|>system<|message|>You are gpt-oss.
reasoning_effort: {{ reasoning_effort }}
<|end|>
"""
PLAIN_TEMPLATE = """
{%- for message in messages %}
{{- message.role + ': ' + message.content + '\\n' }}
{%- endfor %}
"""
# ── Tests: classifier honesty ────────────────────────────────────────
def test_detect_reasoning_flags_qwen3_supports_tools_and_reasoning():
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(QWEN3_TEMPLATE, "unsloth/Qwen3-0.6B")
assert flags["supports_tools"] is True
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking"
assert flags["supports_preserve_thinking"] is True
assert flags["reasoning_always_on"] is False
def test_detect_reasoning_flags_plain_template_all_false():
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(PLAIN_TEMPLATE, "some/PlainChat")
assert flags["supports_tools"] is False
assert flags["supports_reasoning"] is False
assert flags["supports_preserve_thinking"] is False
assert flags["reasoning_always_on"] is False
def test_detect_reasoning_flags_none_template_returns_all_false():
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(None)
assert flags["supports_tools"] is False
assert flags["supports_reasoning"] is False
assert flags["supports_preserve_thinking"] is False
assert flags["reasoning_always_on"] is False
assert flags["reasoning_style"] == "enable_thinking"
def test_detect_safetensors_features_passes_template_through_to_classifier():
"""Route wrapper forwards a real template to the inner classifier."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B")
flags = _detect_safetensors_features(backend, QWEN3_TEMPLATE)
assert flags["supports_tools"] is True
assert flags["supports_reasoning"] is True
def test_detect_safetensors_features_none_template_returns_all_false():
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B")
flags = _detect_safetensors_features(backend, None)
assert flags == {
"supports_reasoning": False,
"reasoning_style": "enable_thinking",
"reasoning_always_on": False,
"supports_preserve_thinking": False,
"supports_tools": False,
}
def test_detect_safetensors_features_gptoss_disables_tools():
"""gpt-oss Harmony: tools intentionally off even if template marks it."""
from routes.inference import _detect_safetensors_features
backend = MagicMock()
backend.active_model_name = "unsloth/gpt-oss-20b"
backend._is_gpt_oss_model.return_value = True
flags = _detect_safetensors_features(backend, QWEN3_TEMPLATE)
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "reasoning_effort"
assert flags["supports_tools"] is False
# Llama-3 / Mistral templates advertise tool handling but the model emits
# tool calls in <|python_tag|> / [TOOL_CALLS] format -- not the
# <tool_call> / <function= our parser understands. The route helper must
# refuse to flip supports_tools=True for those families so the UI does
# not enable a pill the agentic loop cannot honour.
LLAMA3_TEMPLATE = """
{%- if tools %}
{{- '<|start_header_id|>system<|end_header_id|>' }}
{{- 'You have access to the following tools.' }}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
{%- for message in messages %}
{%- if message.role == 'tool' %}
{{- '<|start_header_id|>ipython<|end_header_id|>' }}
{{- '<|python_tag|>' }}
{{- message.content }}
{%- endif %}
{%- endfor %}
"""
MISTRAL_TEMPLATE = """
{%- if tools %}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
{%- for message in messages %}
{%- if message.role == 'tool' %}
{{- '[TOOL_CALLS]' + message.content + '[/TOOL_CALLS]' }}
{%- endif %}
{%- endfor %}
"""
def test_detect_safetensors_features_llama3_template_suppresses_tools():
"""Llama-3 emits <|python_tag|>; safetensors loop cannot parse it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE)
assert flags["supports_tools"] is False
def test_detect_safetensors_features_mistral_template_suppresses_tools():
"""Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3")
flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE)
assert flags["supports_tools"] is False
def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on():
"""Sanity check: gate only suppresses non-Qwen formats."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B")
flags = _detect_safetensors_features(backend, QWEN3_TEMPLATE)
assert flags["supports_tools"] is True
def test_detect_safetensors_features_function_xml_format_keeps_tools_on():
"""Templates emitting <function=name> XML are parser-compatible."""
from routes.inference import _detect_safetensors_features
tpl_with_function_xml = (
"{%- if tools %}<|im_start|>system\n"
"Tool call format: <function=name><parameter=k>v</parameter></function>"
"<|im_end|>{%- endif %}"
)
backend = SimpleNamespace(active_model_name = "custom/with-function-xml")
flags = _detect_safetensors_features(backend, tpl_with_function_xml)
assert flags["supports_tools"] is True
# Qwen3.5 family pins -- the live GGUF + safetensors templates fetched
# from the unsloth/Qwen3.5-0.8B(-GGUF) repos both wrap tool calls as
# ``<tool_call>\n<function=name>...``. Capture a faithful slice so the
# classifier never silently regresses for this family.
QWEN35_TOOL_INSTRUCTION = (
"{%- if tools %}\n"
" <|im_start|>system\n"
" # Tools\n"
" <tools>\n"
" {%- for tool in tools %}{{ tool | tojson }}{%- endfor %}\n"
" </tools>\n"
" If you choose to call a function ONLY reply in the following format:\n"
" <tool_call>\n"
" <function=example_function_name>\n"
" <parameter=example_parameter_1>\n"
" value_1\n"
" </parameter>\n"
" </function>\n"
" </tool_call>\n"
" <|im_end|>\n"
"{%- endif %}\n"
"{%- if enable_thinking is defined and enable_thinking %}{{- '<think>' }}{%- endif %}\n"
)
def test_detect_safetensors_features_qwen35_keeps_tools_on():
"""unsloth/Qwen3.5-0.8B family must surface tools+reasoning enabled."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Qwen3.5-0.8B")
flags = _detect_safetensors_features(backend, QWEN35_TOOL_INSTRUCTION)
assert flags["supports_tools"] is True
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking"
# ── Tests: IPC bridge contract ───────────────────────────────────────
def test_orchestrator_mirrors_chat_template_info_into_models_dict():
"""Worker → orchestrator must copy chat_template_info verbatim."""
from core.inference.orchestrator import InferenceOrchestrator
orch = InferenceOrchestrator.__new__(InferenceOrchestrator)
orch.models = {}
orch.active_model_name = None
orch.loading_models = set()
model_info = {
"identifier": "unsloth/Qwen3-0.6B",
"display_name": "Qwen3-0.6B",
"is_vision": False,
"is_lora": False,
"is_gguf": False,
"is_audio": False,
"audio_type": None,
"has_audio_input": False,
"chat_template_info": {
"has_template": True,
"template": QWEN3_TEMPLATE,
"format_type": "chatml",
"template_name": "qwen3",
"special_tokens": {"bos_token": "<|im_start|>"},
},
}
# Replay orchestrator.load_model's mirror block verbatim.
orch.active_model_name = model_info["identifier"]
orch.models[orch.active_model_name] = {
"is_vision": model_info.get("is_vision", False),
"is_lora": model_info.get("is_lora", False),
"display_name": model_info.get("display_name", "x"),
"is_audio": model_info.get("is_audio", False),
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
}
_tpl_info = model_info.get("chat_template_info")
if isinstance(_tpl_info, dict):
orch.models[orch.active_model_name]["chat_template_info"] = _tpl_info
entry = orch.models[orch.active_model_name]
tpl = entry.get("chat_template_info", {}).get("template")
assert tpl == QWEN3_TEMPLATE
from routes.inference import _detect_safetensors_features
flags = _detect_safetensors_features(
SimpleNamespace(active_model_name = orch.active_model_name), tpl
)
assert flags["supports_tools"] is True
assert flags["supports_reasoning"] is True
def test_orchestrator_missing_chat_template_info_falls_back_to_all_false():
"""Old / malformed worker reply: no crash, all flags False."""
from core.inference.orchestrator import InferenceOrchestrator
from routes.inference import _detect_safetensors_features
orch = InferenceOrchestrator.__new__(InferenceOrchestrator)
orch.models = {}
orch.active_model_name = "unsloth/Qwen3-0.6B"
model_info = {
"identifier": "unsloth/Qwen3-0.6B",
"is_vision": False,
"is_lora": False,
# NB: no chat_template_info key
}
orch.models[orch.active_model_name] = {
"is_vision": False,
"is_lora": False,
}
_tpl_info = model_info.get("chat_template_info")
if isinstance(_tpl_info, dict):
orch.models[orch.active_model_name]["chat_template_info"] = _tpl_info
entry = orch.models[orch.active_model_name]
tpl = entry.get("chat_template_info", {}).get("template")
assert tpl is None
flags = _detect_safetensors_features(
SimpleNamespace(active_model_name = orch.active_model_name), tpl
)
assert flags["supports_tools"] is False
def test_worker_load_reply_payload_includes_chat_template_info():
"""Worker IPC reply carries chat_template_info dict."""
class _StubBackend:
def __init__(self, identifier, template):
self.active_model_name = identifier
self.models = {
identifier: {
"chat_template_info": {
"has_template": True,
"template": template,
"format_type": "chatml",
"template_name": "qwen3",
"special_tokens": {"bos_token": "<|im_start|>"},
}
}
}
backend = _StubBackend("unsloth/Qwen3-0.6B", QWEN3_TEMPLATE)
mc = SimpleNamespace(
identifier = "unsloth/Qwen3-0.6B",
display_name = "Qwen3-0.6B",
is_vision = False,
is_lora = False,
)
# Replay the worker's payload-build block.
model_info = {
"identifier": mc.identifier,
"display_name": mc.display_name,
"is_vision": mc.is_vision,
"is_lora": mc.is_lora,
"is_gguf": False,
}
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
_tpl_info = _entry.get("chat_template_info")
if isinstance(_tpl_info, dict):
model_info["chat_template_info"] = {
"has_template": bool(_tpl_info.get("has_template", False)),
"template": _tpl_info.get("template"),
"format_type": _tpl_info.get("format_type", "generic"),
"template_name": _tpl_info.get("template_name"),
"special_tokens": _tpl_info.get("special_tokens", {}) or {},
}
assert "chat_template_info" in model_info
assert model_info["chat_template_info"]["template"] == QWEN3_TEMPLATE
assert model_info["chat_template_info"]["has_template"] is True
def test_worker_load_reply_payload_survives_missing_template():
"""Tokenizer with no chat_template still produces a valid reply."""
class _StubBackend:
def __init__(self):
self.active_model_name = "legacy/no-template"
self.models = {"legacy/no-template": {}} # no chat_template_info
backend = _StubBackend()
mc = SimpleNamespace(
identifier = "legacy/no-template",
display_name = "legacy",
is_vision = False,
is_lora = False,
)
model_info = {
"identifier": mc.identifier,
"display_name": mc.display_name,
"is_vision": mc.is_vision,
"is_lora": mc.is_lora,
"is_gguf": False,
}
_bm = getattr(backend, "models", {}) or {}
_entry = _bm.get(mc.identifier) or {}
_tpl_info = _entry.get("chat_template_info")
if isinstance(_tpl_info, dict):
model_info["chat_template_info"] = dict(_tpl_info)
assert "chat_template_info" not in model_info
# ── End-to-end: route layer sees the template, advertises True ───────
def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors():
"""End-to-end: Qwen3 safetensors flips supports_tools=True."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(
active_model_name = "unsloth/Qwen3-0.6B",
models = {
"unsloth/Qwen3-0.6B": {
"is_vision": False,
"chat_template_info": {
"has_template": True,
"template": QWEN3_TEMPLATE,
"format_type": "chatml",
},
}
},
)
_model_info = backend.models.get(backend.active_model_name, {})
_tpl = _model_info.get("chat_template_info", {}).get("template")
flags = _detect_safetensors_features(backend, _tpl)
assert flags["supports_tools"] is True
assert flags["supports_reasoning"] is True
assert flags["supports_preserve_thinking"] is True

788
studio/backend/tests/test_safetensors_tool_loop.py Normal file
View file

@ -0,0 +1,788 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Tests for the safetensors agentic tool loop.
Covers the shared ``tool_call_parser`` helpers and the cumulative-text
state machine inside ``safetensors_agentic.run_safetensors_tool_loop``.
The loop is exercised with hand-crafted fake single-turn generators so
no model load is needed; the tests run in CI under a few seconds.
Edge cases under coverage:
* Plain answers (no tool calls) flush full content.
* Single ``<tool_call>{json}</tool_call>`` triggers the tool and re-enters.
* Single ``<function=name>...`` XML form triggers the same path.
* Truncated unclosed ``<tool_call>`` is still parsed.
* Tool result is fed back as ``role=tool`` for the next iteration.
* Bad JSON inside ``<tool_call>`` does not raise and (when healed) is
routed as a ``{"query": ...}`` web search call.
* Duplicate tool calls produce a synthetic "do not repeat" result the
second time.
* ``__IMAGES__`` sentinel is stripped before the model sees the result.
* Tool execution errors are tagged so the model gets a nudge but the
loop keeps streaming.
* Cancel is honoured between iterations.
* ``max_tool_iterations`` cap is respected and a final-answer attempt
closes the stream cleanly.
"""
import threading
import pytest
from core.inference import safetensors_agentic
from core.inference.safetensors_agentic import (
_coerce_arguments,
run_safetensors_tool_loop,
)
from core.inference.tool_call_parser import (
has_tool_signal,
parse_tool_calls_from_text,
strip_tool_markup,
)
from utils.datasets import is_gpt_oss_model_name
# ────────────────────────────────────────────────────────────────────
# parse_tool_calls_from_text
# ────────────────────────────────────────────────────────────────────
class TestParser:
def test_json_tool_call(self):
text = (
'<tool_call>{"name":"web_search","arguments":{"query":"hello"}}</tool_call>'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
tc = result[0]
assert tc["type"] == "function"
assert tc["function"]["name"] == "web_search"
# Arguments must always be a JSON string.
assert isinstance(tc["function"]["arguments"], str)
assert "hello" in tc["function"]["arguments"]
def test_json_tool_call_unclosed(self):
# No </tool_call>; balanced-brace extractor must still close.
text = '<tool_call>{"name":"python","arguments":{"code":"print(1)"}}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
def test_xml_function_call(self):
text = "<function=python><parameter=code>print('hi')</parameter></function>"
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
assert "print('hi')" in result[0]["function"]["arguments"]
def test_xml_unclosed(self):
# Closing tags omitted; parser must still extract the value.
text = "<function=terminal><parameter=command>ls -la"
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "terminal"
assert "ls -la" in result[0]["function"]["arguments"]
def test_code_with_embedded_xml(self):
# A code parameter contains the literal </parameter>. Must not
# truncate the value because the parser uses end-of-body as the
# only boundary for single-parameter calls.
text = (
"<function=python><parameter=code>html = '<a></a>'\n"
"print('hi')</parameter></function>"
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert "print('hi')" in result[0]["function"]["arguments"]
def test_multiple_calls(self):
text = (
'<tool_call>{"name":"web_search","arguments":{"query":"a"}}</tool_call>'
'<tool_call>{"name":"web_search","arguments":{"query":"b"}}</tool_call>'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 2
assert result[0]["function"]["name"] == "web_search"
assert result[1]["function"]["name"] == "web_search"
def test_bad_json_does_not_raise(self):
text = "<tool_call>{not valid json}</tool_call>"
result = parse_tool_calls_from_text(text)
# Bad JSON is silently dropped; caller can fall back to text.
assert result == []
def test_has_tool_signal(self):
assert has_tool_signal("blah <tool_call> x")
assert has_tool_signal("hi <function=foo>...")
assert not has_tool_signal("hello world")
def test_strip_markup_closed(self):
text = "before <tool_call>{}</tool_call> after"
assert strip_tool_markup(text) == "before after"
def test_strip_markup_unclosed_final(self):
text = "before <tool_call>{partial"
# With final=True the trailing run is dropped.
assert strip_tool_markup(text, final = True) == "before"
# Without final=True the unclosed run is preserved.
assert "partial" in strip_tool_markup(text)
# ────────────────────────────────────────────────────────────────────
# run_safetensors_tool_loop
# ────────────────────────────────────────────────────────────────────
def _fake_stream(chunks):
"""Build a single-turn generator that yields cumulative snapshots."""
def _gen(_messages):
acc = ""
for c in chunks:
acc += c
yield acc
return _gen
def _const_stream(text):
"""A single-turn generator that yields one cumulative snapshot."""
def _gen(_messages):
yield text
return _gen
class FakeExecuteTool:
"""Stand-in for ``core.inference.tools.execute_tool``."""
def __init__(self, results):
# ``results`` is a list of strings or RuntimeError instances.
self.results = list(results)
self.calls: list[tuple[str, dict]] = []
def __call__(
self,
name,
arguments,
*,
cancel_event = None,
timeout = None,
session_id = None,
):
self.calls.append((name, arguments))
result = self.results.pop(0) if self.results else "OK"
if isinstance(result, Exception):
raise result
return result
def _collect_events(generator, max_events = 200):
events = []
for ev in generator:
events.append(ev)
if len(events) >= max_events:
break
return events
def _make_loop(*, turns, exec_results = None, **kwargs):
"""Build a configured loop with a multi-turn fake generator.
``turns`` is a list of chunk-lists; iteration N yields chunks from
``turns[N]``.
"""
turn_iter = iter(turns)
def _gen(_messages):
try:
chunks = next(turn_iter)
except StopIteration:
return
acc = ""
for c in chunks:
acc += c
yield acc
exec_fn = FakeExecuteTool(exec_results or [])
return run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "hi"}],
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "python"}},
{"type": "function", "function": {"name": "terminal"}},
],
execute_tool = exec_fn,
**kwargs,
), exec_fn
class TestLoopBasic:
def test_plain_answer(self):
# No tool XML; loop should yield content then status="".
loop, _exec = _make_loop(
turns = [["Hello", " world", "!"]],
exec_results = [],
)
events = _collect_events(loop)
contents = [e for e in events if e["type"] == "content"]
statuses = [e for e in events if e["type"] == "status"]
assert contents, "expected at least one content event"
# Final cumulative content should contain the answer.
final_text = contents[-1]["text"]
assert "Hello world!" in final_text
assert statuses and statuses[-1]["text"] == ""
def test_single_tool_then_answer(self):
loop, exec_fn = _make_loop(
turns = [
# : tool call only.
[
'<tool_call>{"name":"web_search",',
'"arguments":{"query":"weather"}}',
"</tool_call>",
],
# : final answer.
["The ", "weather is ", "sunny."],
],
exec_results = ["Sunny and 22C"],
)
events = _collect_events(loop)
kinds = [e["type"] for e in events]
assert "tool_start" in kinds
assert "tool_end" in kinds
# Tool was actually called with the parsed arguments.
assert exec_fn.calls == [("web_search", {"query": "weather"})]
tool_start = next(e for e in events if e["type"] == "tool_start")
assert tool_start["tool_name"] == "web_search"
tool_end = next(e for e in events if e["type"] == "tool_end")
assert tool_end["result"] == "Sunny and 22C"
contents = [e for e in events if e["type"] == "content"]
assert contents and "sunny" in contents[-1]["text"].lower()
def test_function_xml_form(self):
loop, exec_fn = _make_loop(
turns = [
["<function=python><parameter=code>print(1)</parameter></function>"],
["Result: 1"],
],
exec_results = ["1\n"],
)
events = _collect_events(loop)
assert exec_fn.calls == [("python", {"code": "print(1)"})]
contents = [e for e in events if e["type"] == "content"]
assert "Result: 1" in contents[-1]["text"]
def test_truncated_unclosed_tool_call(self):
loop, exec_fn = _make_loop(
turns = [
# No </tool_call>; balanced-brace parser must still
# succeed because the JSON itself is balanced.
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}'],
["done"],
],
exec_results = ["result"],
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "x"})]
def test_bad_json_healed_to_query(self):
# Tool call with non-JSON string arguments. With auto_heal_tool_calls
# the string is routed as {"query": ...}.
loop, exec_fn = _make_loop(
turns = [
# JSON inside the tool call is well-formed; the
# ``arguments`` is a string that is not itself valid
# JSON for ``_coerce_arguments`` to parse, so the
# heal path runs.
[
'<tool_call>{"name":"web_search","arguments":"hello world"}</tool_call>'
],
["ok"],
],
exec_results = ["..."],
)
events = _collect_events(loop)
assert exec_fn.calls and exec_fn.calls[0][0] == "web_search"
assert exec_fn.calls[0][1] == {"query": "hello world"}
class TestLoopBehaviour:
def test_duplicate_tool_call_synthetic_result(self):
# Two identical successful calls in a row: the second is short-
# circuited with a "do not repeat" message and execute_tool is
# called only once.
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
["final"],
],
exec_results = ["search-result-1"],
)
events = _collect_events(loop)
# Only one real call.
assert len(exec_fn.calls) == 1
tool_end_events = [e for e in events if e["type"] == "tool_end"]
assert len(tool_end_events) == 2
assert "do not repeat" in tool_end_events[1]["result"].lower()
def test_image_sentinel_stripped_from_model_feed(self):
# The tool result has a frontend image sentinel that should be
# stripped before being fed back into the next turn, BUT the
# tool_end event still carries the raw result for the UI.
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"python","arguments":{"code":"plot()"}}</tool_call>'
],
["see chart"],
],
exec_results = ["chart\n__IMAGES__:/tmp/chart.png"],
)
events = _collect_events(loop)
tool_end = next(e for e in events if e["type"] == "tool_end")
assert "__IMAGES__" in tool_end["result"]
def test_image_sentinel_stripped_with_leading_marker(self):
# Sentinel at start (no newline) must not leak to the model.
from core.inference import safetensors_agentic as _sa
captured: list[list[dict]] = []
def fake_single_turn(messages, **_kw):
captured.append([dict(m) for m in messages])
if len(captured) == 1:
yield '<tool_call>{"name":"python","arguments":{"code":"plot()"}}</tool_call>'
else:
yield "done"
events = list(
_sa.run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "plot please"}],
tools = [{"function": {"name": "python"}}],
execute_tool = lambda *_a, **_kw: "__IMAGES__:/tmp/x.png",
cancel_event = threading.Event(),
max_tool_iterations = 3,
auto_heal_tool_calls = True,
)
)
# Model's second turn must not see "__IMAGES__".
assert len(captured) >= 2
tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
assert tool_msgs, "no tool message reached the model"
for tm in tool_msgs:
assert (
"__IMAGES__" not in tm["content"]
), f"sentinel leaked to model: {tm['content']!r}"
def test_image_sentinel_stripped_with_multiple_markers(self):
# Consecutive sentinels: cut at the first, nothing leaks.
from core.inference import safetensors_agentic as _sa
captured: list[list[dict]] = []
def fake_single_turn(messages, **_kw):
captured.append([dict(m) for m in messages])
if len(captured) == 1:
yield '<tool_call>{"name":"python","arguments":{"code":"plot()"}}</tool_call>'
else:
yield "done"
multi = "panel\n__IMAGES__:/tmp/a.png\n__IMAGES__:/tmp/b.png"
events = list(
_sa.run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "plot please"}],
tools = [{"function": {"name": "python"}}],
execute_tool = lambda *_a, **_kw: multi,
cancel_event = threading.Event(),
max_tool_iterations = 3,
auto_heal_tool_calls = True,
)
)
tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
assert tool_msgs
for tm in tool_msgs:
assert (
"__IMAGES__" not in tm["content"]
), f"second sentinel leaked: {tm['content']!r}"
assert (
tm["content"] == "panel"
), f"expected payload-only 'panel', got {tm['content']!r}"
def test_tool_execution_error_is_emitted_but_loop_continues(self):
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
["sorry, that failed"],
],
exec_results = ["Error: network unreachable"],
)
events = _collect_events(loop)
tool_end = next(e for e in events if e["type"] == "tool_end")
assert tool_end["result"].startswith("Error")
# The loop must still produce a content event after the failure.
contents = [e for e in events if e["type"] == "content"]
assert contents
def test_exception_in_executor_does_not_raise(self):
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
["recovered"],
],
exec_results = [RuntimeError("boom")],
)
events = _collect_events(loop)
tool_end = next(e for e in events if e["type"] == "tool_end")
assert "boom" in tool_end["result"]
class TestLoopControl:
def test_cancel_event_breaks_loop(self):
cancel = threading.Event()
cancel.set()
# Even with a fake stream that emits tool calls, the loop must
# bail before invoking execute_tool when cancel is set.
exec_fn = FakeExecuteTool([])
events = list(
run_safetensors_tool_loop(
single_turn = _const_stream(
'<tool_call>{"name":"web_search",'
'"arguments":{"query":"x"}}</tool_call>'
),
messages = [{"role": "user", "content": "hi"}],
tools = [],
execute_tool = exec_fn,
cancel_event = cancel,
)
)
assert events == []
assert exec_fn.calls == []
def test_max_iterations_caps_loop(self):
# The loop should stop after max_tool_iterations even if the
# model keeps asking for tools, then emit a final-attempt round.
loop, exec_fn = _make_loop(
turns = [
# : tool call (executes once)
[
'<tool_call>{"name":"web_search","arguments":{"query":"a"}}</tool_call>'
],
# : model gives a final answer when nudged.
["here is the final answer"],
],
exec_results = ["result"],
max_tool_iterations = 1,
)
events = _collect_events(loop)
contents = [e for e in events if e["type"] == "content"]
# Final content must include the final answer.
assert contents and "final answer" in contents[-1]["text"]
class TestStatusFormatting:
def test_status_for_known_tools(self):
# Use the private helper directly to verify status formatting.
assert (
safetensors_agentic._status_for_tool("web_search", {"query": "abc"})
== "Searching: abc"
)
assert (
safetensors_agentic._status_for_tool(
"web_search", {"url": "https://www.example.com/x"}
)
== "Reading: example.com"
)
assert safetensors_agentic._status_for_tool(
"python", {"code": "x = 1"}
).startswith("Running Python:")
assert safetensors_agentic._status_for_tool(
"terminal", {"command": "ls"}
).startswith("Running:")
assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith(
"Calling:"
)
class TestProseMentioningToolCall:
def test_assistant_prose_with_literal_tool_call_text_survives(self):
# Regression: if the assistant text legitimately mentions
# ``<tool_call>`` as a literal string and the parser finds no
# actual call, the loop must surface the full content instead
# of silently stripping everything past the literal marker.
loop, exec_fn = _make_loop(
turns = [
# : a real tool call so the loop moves to
# .
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
# : prose that mentions the literal text.
["the docs say <tool_call> means an LLM tool call wrapper"],
],
exec_results = ["result"],
)
events = _collect_events(loop)
contents = [e for e in events if e["type"] == "content"]
assert contents, "expected at least one content event"
final = contents[-1]["text"]
assert (
"LLM tool" in final
), f"prose mentioning <tool_call> should not be truncated; got {final!r}"
def test_tool_result_with_tool_call_text_does_not_retrigger(self):
# Tool result text contains the literal ``<tool_call>`` string.
# The loop must only parse the MODEL output, not the tool
# result, so we should see exactly one call.
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
["the docs mention <tool_call> wrappers"],
],
exec_results = ["Page text: <tool_call> appears here in the docs"],
)
events = _collect_events(loop)
assert len(exec_fn.calls) == 1
class TestChatTemplateHelper:
"""Cover the dependency-light helper used by InferenceBackend."""
def setup_method(self):
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
)
self.apply = apply_chat_template_for_generation
class _Tok:
def __init__(self, accepted):
self.accepted = accepted
self.call_count = 0
self.last_kwargs = None
def apply_chat_template(
self, messages, *, tokenize = False, add_generation_prompt = True, **kw
):
self.call_count += 1
unknown = set(kw) - self.accepted
if unknown:
raise TypeError(f"unexpected kwargs: {sorted(unknown)}")
self.last_kwargs = dict(kw)
return "PROMPT"
def test_richest_call_wins_when_template_supports_all(self):
tok = self._Tok({"tools", "enable_thinking"})
self.apply(tok, [], tools = [{}], enable_thinking = True)
assert tok.call_count == 1
assert "tools" in tok.last_kwargs
assert "enable_thinking" in tok.last_kwargs
def test_falls_back_when_template_rejects_reasoning_kwarg(self):
tok = self._Tok({"tools"})
self.apply(tok, [], tools = [{}], enable_thinking = True)
assert tok.call_count >= 2
assert tok.last_kwargs == {"tools": [{}]}
def test_falls_back_to_bare_call(self):
tok = self._Tok(set())
self.apply(tok, [], tools = [{}], enable_thinking = True)
assert tok.last_kwargs == {}
def test_jinja_error_propagates(self):
class Boom:
def apply_chat_template(self, *a, **kw):
raise ValueError("jinja: missing var")
with pytest.raises(ValueError):
self.apply(Boom(), [])
def test_no_kwargs_single_call(self):
tok = self._Tok(set())
self.apply(tok, [])
assert tok.call_count == 1
# ────────────────────────────────────────────────────────────────────
# Guardrails (allowlist, budget, streaming-leak, dedup, id offset,
# auto_heal=False, canonical healed-arg key)
# ────────────────────────────────────────────────────────────────────
class TestGuardrails:
def test_disabled_tool_is_not_executed(self):
exec_fn = FakeExecuteTool([])
loop = run_safetensors_tool_loop(
single_turn = _fake_stream(
[
'<tool_call>{"name":"terminal","arguments":{"command":"echo bypass"}}</tool_call>'
]
),
messages = [{"role": "user", "content": "hi"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 2,
)
events = _collect_events(loop)
assert exec_fn.calls == []
tool_ends = [e for e in events if e["type"] == "tool_end"]
assert tool_ends and "not enabled" in tool_ends[0]["result"].lower()
def test_empty_tools_list_does_not_enforce_allowlist(self):
exec_fn = FakeExecuteTool(["OK"])
loop = run_safetensors_tool_loop(
single_turn = _fake_stream(
[
'<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>'
]
),
messages = [{"role": "user", "content": "hi"}],
tools = [],
execute_tool = exec_fn,
max_tool_iterations = 2,
)
_collect_events(loop)
assert exec_fn.calls == [("python", {"code": "print(1)"})]
def test_max_iterations_zero_executes_no_tools(self):
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
]
],
exec_results = ["OK"],
max_tool_iterations = 0,
)
events = _collect_events(loop)
assert exec_fn.calls == []
assert events and events[-1] == {"type": "status", "text": ""}
def test_streaming_clips_before_tool_signal_no_leak(self):
loop, exec_fn = _make_loop(
turns = [
[
"I will look this up. ",
"Some more prose that's long enough to leave the buffer. ",
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>',
],
["all done"],
],
exec_results = ["weather: sunny"],
max_tool_iterations = 2,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "x"})]
for e in events:
if e["type"] == "content":
assert "<tool_call>" not in e["text"]
assert "web_search" not in e["text"]
def test_auto_heal_disabled_still_parses_valid_tool_call(self):
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
],
["done"],
],
exec_results = ["OK"],
auto_heal_tool_calls = False,
max_tool_iterations = 2,
)
_collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "x"})]
def test_non_consecutive_duplicate_is_short_circuited(self):
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"A"}}</tool_call>'
],
[
'<tool_call>{"name":"web_search","arguments":{"query":"B"}}</tool_call>'
],
[
'<tool_call>{"name":"web_search","arguments":{"query":"A"}}</tool_call>'
],
["final"],
],
exec_results = ["res-A", "res-B"],
max_tool_iterations = 4,
)
events = _collect_events(loop)
assert exec_fn.calls == [
("web_search", {"query": "A"}),
("web_search", {"query": "B"}),
]
tool_ends = [e for e in events if e["type"] == "tool_end"]
assert "already made this exact call" in tool_ends[-1]["result"]
def test_coerce_string_args_python_uses_code_key(self):
assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {
"code": "print(1)"
}
def test_coerce_string_args_terminal_uses_command_key(self):
assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == {
"command": "ls -la"
}
def test_tool_call_ids_unique_across_loop_iterations(self):
loop, _exec = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"A"}}</tool_call>'
],
[
'<tool_call>{"name":"web_search","arguments":{"query":"B"}}</tool_call>'
],
["done"],
],
exec_results = ["A", "B"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
ids = [e["tool_call_id"] for e in events if e["type"] == "tool_start"]
assert len(ids) == 2 and ids[0] != ids[1]
# ────────────────────────────────────────────────────────────────────
# Shared gpt-oss name detector
# ────────────────────────────────────────────────────────────────────
class TestGptOssNameDetection:
def test_substring_match(self):
assert is_gpt_oss_model_name("unsloth/gpt-oss-20b") is True
def test_negative_known_non_oss_model(self):
assert is_gpt_oss_model_name("meta-llama/Llama-3.1-8B-Instruct") is False
def test_empty_or_none_returns_false(self):
assert is_gpt_oss_model_name("") is False
assert is_gpt_oss_model_name(None) is False
if __name__ == "__main__":
pytest.main([__file__, "-v"])

2
studio/backend/utils/datasets/__init__.py
View file

@ -59,6 +59,7 @@ from .model_mappings import (
TEMPLATE_TO_MODEL_MAPPER,
MODEL_TO_TEMPLATE_MAPPER,
TEMPLATE_TO_RESPONSES_MAPPER,
is_gpt_oss_model_name,
)
# Legacy imports from the original dataset_utils.py for backward compatibility
@ -98,6 +99,7 @@ __all__ = [
"TEMPLATE_TO_MODEL_MAPPER",
"MODEL_TO_TEMPLATE_MAPPER",
"TEMPLATE_TO_RESPONSES_MAPPER",
"is_gpt_oss_model_name",
# Main entry points
"check_dataset_format",
"format_and_template_dataset",

20
studio/backend/utils/datasets/model_mappings.py
View file

@ -442,6 +442,26 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items():
MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key
def is_gpt_oss_model_name(name: str) -> bool:
"""Name-based check for gpt-oss / harmony models.
Used by both the in-process backend and the parent-process
orchestrator to detect harmony models without an IPC round-trip.
"""
name = (name or "").lower()
if not name:
return False
try:
if MODEL_TO_TEMPLATE_MAPPER.get(name) == "gpt-oss":
return True
for key, tmpl in MODEL_TO_TEMPLATE_MAPPER.items():
if tmpl == "gpt-oss" and (key in name or name in key):
return True
except Exception:
pass
return "gpt-oss" in name
TEMPLATE_TO_RESPONSES_MAPPER = {
"gemma-4-thinking": {
"instruction": "<|turn>user\n",

10
studio/frontend/src/app/auth-guards.ts
View file

@ -9,6 +9,7 @@ import {
hasRefreshToken,
mustChangePassword,
refreshSession,
setMustChangePassword,
} from "@/features/auth";
async function hasActiveSession(): Promise<boolean> {
@ -26,7 +27,12 @@ async function fetchAuthStatus(): Promise<AuthStatus> {
try {
const res = await fetch(apiUrl("/api/auth/status"));
if (!res.ok) return { initialized: true, requires_password_change: mustChangePassword() };
return (await res.json()) as AuthStatus;
const status = (await res.json()) as AuthStatus;
// Server truth wins; keep localStorage in sync both ways.
if (status.requires_password_change !== mustChangePassword()) {
setMustChangePassword(status.requires_password_change);
}
return status;
} catch {
return { initialized: true, requires_password_change: mustChangePassword() };
}
@ -61,6 +67,8 @@ export async function requireGuest(): Promise<void> {
throw redirect({ to: "/chat" });
}
if (!(await hasActiveSession())) return;
// Reconcile localStorage before routing.
await fetchAuthStatus();
throw redirect({ to: getPostAuthRoute() });
}

14
studio/frontend/src/app/routes/__root.tsx
View file

@ -15,9 +15,17 @@ import {
useRouterState,
} from "@tanstack/react-router";
import { AnimatePresence, motion } from "motion/react";
import { Suspense, useEffect } from "react";
import { Suspense, useEffect, type ReactNode } from "react";
import { AppProvider } from "../provider";
// Fallback while a lazy route bundle (Train/Recipes/Export) loads.
// /chat is synchronous and never hits this.
const RouteFallback: ReactNode = (
<div className="flex h-full min-h-0 flex-1 items-center justify-center text-muted-foreground text-sm">
Loading...
</div>
);
const CHAT_ONLY_ALLOWED = new Set([
"/",
"/chat",
@ -72,7 +80,7 @@ function RootLayout() {
<SettingsDialog />
{hideNavbar ? (
<main className="flex-1">
<Suspense fallback={null}>
<Suspense fallback={RouteFallback}>
<Outlet />
</Suspense>
</main>
@ -98,7 +106,7 @@ function RootLayout() {
transition={{ duration: 0.15 }}
className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"}`}
>
<Suspense fallback={null}>
<Suspense fallback={RouteFallback}>
<Outlet />
</Suspense>
</motion.div>

11
studio/frontend/src/components/assistant-ui/attachment.tsx
View file

@ -120,12 +120,13 @@ const AttachmentPreviewDialog: FC<PropsWithChildren> = ({ children }) => {
const AttachmentThumb: FC = () => {
const src = useAttachmentSrc();
const name = useAuiState(({ attachment }) => attachment.name);
if (src) {
return (
<img
src={src}
alt="Attachment preview"
alt={name || "Attachment preview"}
className="h-full w-full object-cover"
/>
);
@ -143,6 +144,7 @@ const AttachmentUI: FC = () => {
const isComposer = aui.attachment.source === "composer";
const isImage = useAuiState(({ attachment }) => attachment.type === "image");
const name = useAuiState(({ attachment }) => attachment.name);
const typeLabel = useAuiState(({ attachment }) => {
const type = attachment.type;
switch (type) {
@ -156,6 +158,11 @@ const AttachmentUI: FC = () => {
throw new Error(`Unknown attachment type: ${type as string}`);
}
});
// Include filename in accessible name so screen readers distinguish
// same-typed attachments. Sighted users get it via the tooltip.
const accessibleName = name
? `${typeLabel} attachment: ${name}`
: `${typeLabel} attachment`;
return (
<Tooltip>
@ -175,7 +182,7 @@ const AttachmentUI: FC = () => {
"aui-attachment-tile-composer border-foreground/20",
)}
id="attachment-tile"
aria-label={`${typeLabel} attachment`}
aria-label={accessibleName}
type="button"
>
<AttachmentThumb />

23
studio/frontend/src/components/assistant-ui/message-timing.tsx
View file

@ -38,9 +38,22 @@ export const MessageTiming: FC<{
)?.custom as { serverTimings?: Record<string, number> } | undefined;
const st = serverTimings?.serverTimings;
// 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
// non-zero decode window AND a finite rate. Fast cached single-token
// responses (sub-10ms) are legitimate and must stay visible.
const hasPredicted =
(st?.predicted_n ?? 0) >= 1 && (st?.predicted_ms ?? 0) > 0;
const predictedRate =
hasPredicted &&
st?.predicted_per_second != null &&
Number.isFinite(st.predicted_per_second)
? st.predicted_per_second
: undefined;
// Badge text: show tok/s if available, otherwise total time
const badgeText = st?.predicted_per_second != null
? `${st.predicted_per_second.toFixed(1)} tok/s`
const badgeText = predictedRate != null
? `${predictedRate.toFixed(1)} tok/s`
: formatTimingMs(timing.totalStreamTime);
return (
@ -85,7 +98,7 @@ export const MessageTiming: FC<{
</span>
</div>
)}
{st?.predicted_ms != null && (
{hasPredicted && st?.predicted_ms != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Generation</span>
<span className="font-mono tabular-nums">
@ -93,11 +106,11 @@ export const MessageTiming: FC<{
</span>
</div>
)}
{st?.predicted_per_second != null && (
{predictedRate != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Speed</span>
<span className="font-mono tabular-nums">
{st.predicted_per_second.toFixed(1)} tok/s
{predictedRate.toFixed(1)} tok/s
</span>
</div>
)}

21
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -916,6 +916,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
<ComposerPrimitive.Dictate asChild={true}>
<TooltipIconButton
tooltip="Dictate"
aria-label="Dictate"
variant="ghost"
className="size-8 rounded-full text-muted-foreground"
>
@ -927,6 +928,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
<ComposerPrimitive.StopDictation asChild={true}>
<TooltipIconButton
tooltip="Stop dictation"
aria-label="Stop dictation"
variant="ghost"
className="size-8 rounded-full text-destructive"
>
@ -994,6 +996,24 @@ const GeneratingIndicator: FC = () => {
return <span className="text-sm text-muted-foreground">Generating...</span>;
};
// Placeholder when stop fires before any visible content (e.g. mid-think).
const CancelledIndicator: FC = () => {
const show = useAuiState(
({ message }) =>
message.content.length === 0 &&
message.status?.type === "incomplete" &&
message.status?.reason === "cancelled",
);
if (!show) {
return null;
}
return (
<span className="aui-cancelled-indicator text-sm italic text-muted-foreground">
Cancelled.
</span>
);
};
const AssistantMessage: FC = () => {
return (
<MessagePrimitive.Root
@ -1002,6 +1022,7 @@ const AssistantMessage: FC = () => {
>
<div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
<GeneratingIndicator />
<CancelledIndicator />
<MessagePrimitive.Parts
components={{
Text: MarkdownText,

14
studio/frontend/src/features/auth/api.ts
View file

@ -65,7 +65,11 @@ async function redirectToAuth(): Promise<void> {
const res = await fetch(apiUrl("/api/auth/status"));
if (res.ok) {
const data = (await res.json()) as { requires_password_change: boolean };
if (data.requires_password_change || mustChangePassword()) target = "/change-password";
// Server truth wins; keep localStorage in sync both ways.
if (data.requires_password_change !== mustChangePassword()) {
setMustChangePassword(data.requires_password_change);
}
if (data.requires_password_change) target = "/change-password";
}
} catch {
// Fall through to /login on error
@ -155,6 +159,14 @@ export async function authFetch(
});
} catch (err) {
if (err instanceof TypeError) {
// fetch TypeError = offline | backend down | CORS/DNS. In Tauri
// it's always backend-down; in the web build distinguish offline
// so the user gets the right recovery path.
if (!isTauri && typeof navigator !== "undefined" && navigator.onLine === false) {
throw new Error(
"You appear to be offline. Check your network connection and try again.",
);
}
throw new Error("Studio isn't running -- please relaunch it.");
}
throw err;

11
studio/frontend/src/features/auth/components/auth-form.tsx
View file

@ -105,12 +105,17 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
setInitialized(result.initialized);
setRequiresPasswordChange(result.requires_password_change);
// Server truth wins; keep localStorage in sync both ways.
if (result.requires_password_change !== mustChangePassword()) {
setMustChangePassword(result.requires_password_change);
}
// Redirect between login ↔ change-password based on server state
if (mode === "login" && result.requires_password_change) {
navigate({ to: "/change-password" });
return;
}
if (mode === "change-password" && !result.requires_password_change && !mustChangePassword()) {
if (mode === "change-password" && !result.requires_password_change) {
navigate({ to: "/login" });
return;
}
@ -163,14 +168,14 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
const blockedByState =
initialized === false ||
(mode === "login" && requiresPasswordChange) ||
(mode === "change-password" && !requiresPasswordChange && !mustChangePassword());
(mode === "change-password" && !requiresPasswordChange);
let helperText: string | null = null;
if (initialized === false) {
helperText = "Auth is still bootstrapping the default admin account.";
} else if (isLoginMode && requiresPasswordChange) {
helperText = "Sign in once with the seeded credentials to change the password.";
} else if (!isLoginMode && !requiresPasswordChange && !mustChangePassword()) {
} else if (!isLoginMode && !requiresPasswordChange) {
helperText = "Password already updated. Use the login screen.";
}
const title = isLoginMode ? "Welcome back" : "Setup your account";

26
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -437,6 +437,9 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise<void> {
* without selecting one. Prefers GGUF (picks smallest cached variant),
* falls back to smallest cached safetensors model.
*/
// Cap cascade so broken cached repos can't spam /api/inference/load.
const MAX_AUTO_LOAD_ATTEMPTS = 3;
async function autoLoadSmallestModel(): Promise<{
loaded: boolean;
blockedByTrustRemoteCode: boolean;
@ -451,6 +454,7 @@ async function autoLoadSmallestModel(): Promise<{
});
let blockedByTrustRemoteCode = false;
let hadNonTrustFailure = false;
let loadAttempts = 0;
async function canAutoLoad(payload: {
model_path: string;
@ -481,6 +485,7 @@ async function autoLoadSmallestModel(): Promise<{
if (ggufRepos.length > 0) {
const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes);
for (const repo of sorted) {
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
try {
const variants = await listGgufVariants(repo.repo_id);
const downloaded = variants.variants
@ -498,6 +503,7 @@ async function autoLoadSmallestModel(): Promise<{
) {
continue;
}
loadAttempts += 1;
const loadResp = await loadModel({
model_path: repo.repo_id,
hf_token: hfToken,
@ -560,6 +566,7 @@ async function autoLoadSmallestModel(): Promise<{
if (modelRepos.length > 0) {
const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes);
for (const repo of sorted) {
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
try {
if (
!(await canAutoLoad({
@ -571,6 +578,7 @@ async function autoLoadSmallestModel(): Promise<{
) {
continue;
}
loadAttempts += 1;
const sfLoadResp = await loadModel({
model_path: repo.repo_id,
hf_token: hfToken,
@ -593,6 +601,12 @@ async function autoLoadSmallestModel(): Promise<{
reasoningStyle: sfLoadResp.reasoning_style ?? "enable_thinking",
supportsPreserveThinking: sfLoadResp.supports_preserve_thinking ?? false,
supportsTools: sfLoadResp.supports_tools ?? false,
// Parity with the GGUF branch above.
toolsEnabled: sfLoadResp.supports_tools ?? false,
codeToolsEnabled: sfLoadResp.supports_tools ?? false,
defaultChatTemplate: sfLoadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
});
const sfModel: ChatModelSummary = {
id: repo.repo_id,
@ -616,6 +630,17 @@ async function autoLoadSmallestModel(): Promise<{
}
}
// Cap also gates the default download so the total /api/inference/load
// budget across cached + fallback is MAX_AUTO_LOAD_ATTEMPTS, not +1.
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
toast.dismiss(toastId);
return {
loaded: false,
blockedByTrustRemoteCode:
blockedByTrustRemoteCode && !hadNonTrustFailure,
};
}
// No cached models found — try downloading a small default GGUF
toast("Downloading a small model…", {
id: toastId,
@ -634,6 +659,7 @@ async function autoLoadSmallestModel(): Promise<{
toast.dismiss(toastId);
return { loaded: false, blockedByTrustRemoteCode };
}
loadAttempts += 1;
const loadResp = await loadModel({
model_path: "unsloth/gemma-4-E2B-it-GGUF",
hf_token: hfToken,

91
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -570,6 +570,11 @@ export function ChatSettingsPanel({
const loadedSpeculativeType = useChatRuntimeStore(
(s) => s.loadedSpeculativeType,
);
const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax);
const setSpecDraftNMax = useChatRuntimeStore((s) => s.setSpecDraftNMax);
const loadedSpecDraftNMax = useChatRuntimeStore(
(s) => s.loadedSpecDraftNMax,
);
const modelRequiresTrustRemoteCode = useChatRuntimeStore(
(s) => s.modelRequiresTrustRemoteCode,
);
@ -608,7 +613,8 @@ export function ChatSettingsPanel({
const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
const ctxDirty = customContextLength !== null;
const specDirty = speculativeType !== loadedSpeculativeType;
const modelSettingsDirty = kvDirty || ctxDirty || specDirty;
const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax;
const modelSettingsDirty = kvDirty || ctxDirty || specDirty || specDraftDirty;
const chatTemplateOverride = useChatRuntimeStore(
(s) => s.chatTemplateOverride,
);
@ -985,19 +991,81 @@ export function ChatSettingsPanel({
Speculative Decoding
</span>
<InfoHint>
Faster generation with 0% accuracy hit.
Faster generation with 0% accuracy hit. Auto picks
MTP / ngram-mod based on the model and platform.
Pick MTP, Ngram, or MTP+Ngram to force a specific
strategy on both GPU and CPU.
</InfoHint>
</div>
<Switch
className="panel-switch shrink-0"
checked={
speculativeType !== "off" && speculativeType != null
}
onCheckedChange={(checked) => {
setSpeculativeType(checked ? "default" : "off");
}}
/>
<div className="flex shrink-0 items-center gap-1.5">
<Select
value={speculativeType ?? "auto"}
onValueChange={(v) => {
setSpeculativeType(v);
if (v !== "mtp" && v !== "mtp+ngram") {
setSpecDraftNMax(null);
}
}}
>
<SelectTrigger
animateRadius={false}
icon={ArrowDown01Icon}
iconClassName="size-3.5"
className="grid h-7 w-[120px] min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 rounded-[10px] border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.07] px-2 py-0 text-[13px]! font-medium text-nav-fg focus-visible:ring-0 focus-visible:border-transparent [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0"
data-test-id="speculative-type-select"
>
<SelectValue />
</SelectTrigger>
<SelectContent className="menu-soft-surface ring-0 border-0 rounded-lg">
<SelectItem value="auto">Auto</SelectItem>
<SelectItem value="mtp">MTP</SelectItem>
<SelectItem value="ngram">Ngram</SelectItem>
<SelectItem value="mtp+ngram">MTP+Ngram</SelectItem>
<SelectItem value="off">Off</SelectItem>
</SelectContent>
</Select>
</div>
</div>
{(speculativeType === "mtp" ||
speculativeType === "mtp+ngram") && (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Draft Tokens
</span>
<InfoHint>
Max MTP draft tokens per step
(--spec-draft-n-max). Lower = less wasted
draft decode; higher = bigger speedup when
acceptance stays high. Default: 2 on GPU,
3 on CPU/Mac.
</InfoHint>
</div>
<input
type="number"
min={1}
max={16}
step={1}
value={specDraftNMax ?? ""}
placeholder="auto"
onChange={(e) => {
const raw = e.target.value;
if (raw === "") {
setSpecDraftNMax(null);
return;
}
const parsed = Number.parseInt(raw, 10);
if (Number.isFinite(parsed)) {
const clamped = Math.max(1, Math.min(16, parsed));
setSpecDraftNMax(clamped);
}
}}
data-test-id="spec-draft-n-max-input"
aria-label="Speculative decoding draft tokens"
className="h-7 w-[72px] rounded-[10px] border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.07] px-2 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0"
/>
</div>
)}
</>
)}
{!isGguf && params.checkpoint && (
@ -1051,6 +1119,7 @@ export function ChatSettingsPanel({
setCustomContextLength(null);
setKvCacheDtype(loadedKvCacheDtype);
setSpeculativeType(loadedSpeculativeType);
setSpecDraftNMax(loadedSpecDraftNMax);
setChatTemplateOverride(loadedChatTemplateOverride);
}}
className="h-7 px-3 text-[12px] font-medium tracking-nav text-muted-foreground"

52
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -141,10 +141,30 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string {
return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`;
}
// Canonicalises any value the backend reports (or persisted state holds)
// onto the five UI-facing modes the Speculative Decoding dropdown
// understands: "auto" / "mtp" / "ngram" / "mtp+ngram" / "off" / null.
// Mirrors backend _canonicalize_spec_mode so old persisted "default" /
// "draft-mtp" / "ngram-mod" / chain values round-trip cleanly.
function normalizeSpeculativeType(v: string | null | undefined): string | null {
if (v == null) return null;
if (v === "default" || v === "off") return v;
return "default";
const s = String(v).trim().toLowerCase();
if (!s) return null;
if (s === "auto" || s === "default") return "auto";
if (s === "off") return "off";
if (s === "ngram-simple") return "ngram-simple";
if (s === "mtp" || s === "draft-mtp") return "mtp";
if (s === "ngram" || s === "ngram-mod") return "ngram";
if (s === "mtp+ngram") return "mtp+ngram";
// Comma-chained legacy values (e.g. from older persisted state).
const parts = s.split(",").map((p) => p.trim()).filter(Boolean);
const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp");
const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod");
if (hasMtp && hasNgram) return "mtp+ngram";
if (hasMtp) return "mtp";
if (hasNgram) return "ngram";
// Unknown -> safe fallback to Auto so the dropdown stays controlled.
return "auto";
}
type LocalReasoningEffort = Extract<ReasoningEffort, "low" | "medium" | "high">;
@ -323,6 +343,12 @@ export function useChatModelRuntime() {
speculativeType: currentSpecType,
loadedSpeculativeType: currentSpecType,
}),
...(statusRes.spec_draft_n_max !== undefined &&
prevState.loadedSpecDraftNMax === null &&
prevState.specDraftNMax === null && {
specDraftNMax: statusRes.spec_draft_n_max ?? null,
loadedSpecDraftNMax: statusRes.spec_draft_n_max ?? null,
}),
...(statusRes.cache_type_kv !== undefined &&
prevState.loadedKvCacheDtype === null && {
kvCacheDtype: statusRes.cache_type_kv,
@ -528,12 +554,31 @@ export function useChatModelRuntime() {
}
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
// Reset Speculative Decoding to Auto whenever the user
// switches to a different model. Spec strategy is a
// per-model decision: a sub-3B non-MTP GGUF that ran with
// "Off" should not carry that choice into a 27B MTP GGUF
// where Auto would auto-promote to draft-mtp. The user can
// still pick a forced mode on the new model; this just
// clears the stale prior-model choice so the backend's
// platform-aware path runs by default. Same applies to
// spec_draft_n_max which is MTP-only.
if (currentCheckpoint && currentCheckpoint !== modelId) {
useChatRuntimeStore.setState({
speculativeType: null,
loadedSpeculativeType: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
});
}
const {
chatTemplateOverride,
kvCacheDtype,
customContextLength,
ggufContextLength,
speculativeType,
specDraftNMax,
activePresetSource,
activeGgufVariant,
} = useChatRuntimeStore.getState();
@ -561,6 +606,7 @@ export function useChatModelRuntime() {
chat_template_override: effectiveChatTemplateOverride,
cache_type_kv: kvCacheDtype,
speculative_type: speculativeType,
spec_draft_n_max: specDraftNMax,
});
// If cancelled while loading, don't update UI to show
@ -635,6 +681,8 @@ export function useChatModelRuntime() {
loadedKvCacheDtype: loadedKv,
speculativeType: loadedSpec,
loadedSpeculativeType: loadedSpec,
specDraftNMax: loadResponse.spec_draft_n_max ?? null,
loadedSpecDraftNMax: loadResponse.spec_draft_n_max ?? null,
customContextLength: keepCustomCtx,
defaultChatTemplate: loadResponse.chat_template ?? null,
chatTemplateOverride: effectiveChatTemplateOverride,

14
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -523,7 +523,19 @@ export function SharedComposer({
handlesRef.current["model1"] || handlesRef.current["model2"],
);
const isGeneralizedCompare =
hasCompareHandles && Boolean(model1?.id || model2?.id);
hasCompareHandles && Boolean(model1?.id && model2?.id);
// Generalized compare requires both panes to have a model. A
// half-selected send either races to an empty bubble with bogus
// tok/s (#5569) or leaves the empty pane with a dangling prompt.
// hasCompareHandles is true only in GeneralCompareContent, so
// LoraCompare and single-pane chats are unaffected.
if (hasCompareHandles && !isGeneralizedCompare) {
toast.error("Pick a model in each pane to compare", {
description: "Use the model dropdown above each pane, then send your prompt.",
});
return;
}
if (pendingImages.length > 0 && !isGeneralizedCompare && imageUnavailableReason) {
// Single mode: the loaded model's runtime capability is known

13
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -258,6 +258,9 @@ type ChatRuntimeStore = {
loadedKvCacheDtype: string | null;
speculativeType: string | null;
loadedSpeculativeType: string | null;
/** User --spec-draft-n-max override (null = platform default). */
specDraftNMax: number | null;
loadedSpecDraftNMax: number | null;
loadedIsMultimodal: boolean;
customContextLength: number | null;
defaultChatTemplate: string | null;
@ -305,6 +308,7 @@ type ChatRuntimeStore = {
setToolCallTimeout: (value: number) => void;
setKvCacheDtype: (dtype: string | null) => void;
setSpeculativeType: (type: string | null) => void;
setSpecDraftNMax: (value: number | null) => void;
setCustomContextLength: (v: number | null) => void;
setChatTemplateOverride: (template: string | null) => void;
setPendingAudio: (base64: string, name: string) => void;
@ -349,8 +353,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5),
kvCacheDtype: null,
loadedKvCacheDtype: null,
speculativeType: "default",
speculativeType: "auto",
loadedSpeculativeType: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
customContextLength: null,
defaultChatTemplate: null,
@ -457,8 +463,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
toolStatus: null,
kvCacheDtype: null,
loadedKvCacheDtype: null,
speculativeType: "default",
speculativeType: "auto",
loadedSpeculativeType: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
customContextLength: null,
defaultChatTemplate: null,
@ -506,6 +514,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
}),
setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
setSpeculativeType: (speculativeType) => set({ speculativeType }),
setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }),
setCustomContextLength: (customContextLength) => set({ customContextLength }),
setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }),
setPendingAudio: (base64, name) =>

18
studio/frontend/src/features/chat/types/api.ts
View file

@ -42,7 +42,21 @@ export interface LoadModelRequest {
trust_remote_code?: boolean;
chat_template_override?: string | null;
cache_type_kv?: string | null;
/**
* Speculative decoding mode for GGUF models. Canonical values:
* "auto" (platform-aware: MTP on MTP GGUFs, ngram-mod fallback for
* sub-3B), "mtp" (force draft-mtp only on both GPU and CPU), "ngram"
* (force ngram-mod only), "mtp+ngram" (force ngram-mod + draft-mtp
* chain on both platforms), or "off". Legacy values "default" /
* "draft-mtp" / "ngram-mod" / "ngram-simple" are still accepted by
* the backend.
*/
speculative_type?: string | null;
/**
* Override --spec-draft-n-max for MTP speculative decoding. Only
* applied when speculative_type resolves to "mtp" or "mtp+ngram".
*/
spec_draft_n_max?: number | null;
}
export interface ValidateModelResponse {
@ -118,7 +132,9 @@ export interface LoadModelResponse {
supports_tools?: boolean;
cache_type_kv?: string | null;
chat_template?: string | null;
/** Canonical UI-facing mode the load request resolved to. See LoadModelRequest. */
speculative_type?: string | null;
spec_draft_n_max?: number | null;
}
export interface UnloadModelRequest {
@ -155,7 +171,9 @@ export interface InferenceStatusResponse {
native_context_length?: number | null;
cache_type_kv?: string | null;
chat_template_override?: string | null;
/** Canonical UI-facing mode currently active. See LoadModelRequest. */
speculative_type?: string | null;
spec_draft_n_max?: number | null;
}
export interface AudioGenerationResponse {

20
studio/frontend/src/features/settings/settings-dialog.tsx
View file

@ -74,6 +74,7 @@ export function SettingsDialog() {
const activeTab = useSettingsDialogStore((s) => s.activeTab);
const setActiveTab = useSettingsDialogStore((s) => s.setActiveTab);
const closeDialog = useSettingsDialogStore((s) => s.closeDialog);
const opener = useSettingsDialogStore((s) => s.opener);
const reduced = useReducedMotion();
const tabButtonRefs = useRef<Record<SettingsTab, HTMLButtonElement | null>>({
general: null,
@ -98,11 +99,22 @@ export function SettingsDialog() {
<DialogContent
showCloseButton={false}
overlayClassName="bg-background/40"
onCloseAutoFocus={(e) => {
// Restore focus to the element that triggered openDialog().
// Radix's FocusScope races our rAF-scheduled tab-button focus
// and loses the previous-focus reference, so we restore by hand.
if (opener && opener.isConnected) {
e.preventDefault();
opener.focus({ preventScroll: true });
}
}}
className={cn(
"!max-w-none h-[560px] w-[820px] p-0 overflow-hidden",
// Cap at 820px but shrink to the viewport so we don't clip
// on iPad-portrait widths (640-820px) where the fixed
// `w-[820px]` overflows by 26px on each side.
"!max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden",
"shadow-border rounded-xl border-border",
"sm:h-[560px] sm:w-[820px]",
"max-sm:h-dvh max-sm:w-dvw max-sm:rounded-none",
"max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none",
)}
>
<DialogTitle className="sr-only">Settings</DialogTitle>
@ -110,7 +122,7 @@ export function SettingsDialog() {
Manage your Unsloth Studio preferences.
</DialogDescription>
<div className="flex h-full min-h-0 max-sm:flex-col">
<aside className="font-heading flex w-[200px] shrink-0 flex-col border-r border-border bg-muted/20 p-2 max-sm:w-full max-sm:border-r-0 max-sm:border-b">
<aside className="font-heading flex w-[216px] shrink-0 flex-col border-r border-border bg-muted/20 p-2 max-sm:w-full max-sm:border-r-0 max-sm:border-b">
<nav className="flex flex-col gap-0.5 max-sm:flex-row max-sm:overflow-x-auto">
{TABS.map((tab) => {
const active = activeTab === tab.id;

15
studio/frontend/src/features/settings/stores/settings-dialog-store.ts
View file

@ -15,6 +15,11 @@ export type SettingsTab =
interface SettingsDialogState {
open: boolean;
activeTab: SettingsTab;
// Element focused at the moment openDialog() ran. Radix's FocusScope
// would normally track this, but the rAF-scheduled focus() in
// settings-dialog.tsx races its previous-focus capture, leaving focus
// on <body> after close. We restore explicitly via onCloseAutoFocus.
opener: HTMLElement | null;
openDialog: (tab?: SettingsTab) => void;
closeDialog: () => void;
setActiveTab: (tab: SettingsTab) => void;
@ -47,11 +52,21 @@ function loadInitialTab(): SettingsTab {
export const useSettingsDialogStore = create<SettingsDialogState>((set) => ({
open: false,
activeTab: loadInitialTab(),
opener: null,
openDialog: (tab) =>
set((state) => ({
open: true,
activeTab: tab ?? state.activeTab,
opener:
typeof document !== "undefined" &&
document.activeElement instanceof HTMLElement &&
document.activeElement !== document.body
? document.activeElement
: null,
})),
// Do NOT clear `opener` here. onCloseAutoFocus runs on the next render
// pass after `open: false` lands, so the opener must still be readable
// from the store at that point. The next openDialog() overwrites it.
closeDialog: () => set({ open: false }),
setActiveTab: (tab) => {
try {

7
studio/frontend/src/features/settings/stores/theme-store.ts
View file

@ -32,7 +32,12 @@ function resolveTheme(theme: Theme): ResolvedTheme {
function applyToDocument(resolved: ResolvedTheme) {
if (typeof document === "undefined") return;
document.documentElement.classList.toggle("dark", resolved === "dark");
// Keep "dark"/"light" mutually exclusive. next-themes (via Sonner)
// adds "light" on first mount; without the explicit toggle we'd end
// up with `class="light dark"` after a switch.
const cl = document.documentElement.classList;
cl.toggle("dark", resolved === "dark");
cl.toggle("light", resolved === "light");
}
const listeners = new Set<() => void>();

42
studio/frontend/src/index.css
View file

@ -1172,11 +1172,49 @@
mix-blend-mode: normal;
}
/* Override sonner's hardcoded top: 0 on the toast close button. */
/* Override sonner top: 0 and pin to theme tokens (--gray2 hover ignores data-sonner-theme). */
[data-sonner-toast][data-styled="true"] [data-close-button] {
top: 8px !important;
background: var(--popover) !important;
color: var(--popover-foreground) !important;
border-color: var(--border) !important;
}
/* Bump the X stroke so it stays visible against dark backgrounds. */
[data-sonner-toast][data-styled="true"] [data-close-button] svg {
stroke-width: 2.25;
}
[data-sonner-toast][data-styled="true"]:hover [data-close-button]:hover {
background: var(--muted) !important;
color: var(--popover-foreground) !important;
border-color: var(--border) !important;
}
/*
* prefers-reduced-motion: honour the OS-level "reduce motion" preference.
* Tailwind animate-in/out, Radix open/close transforms, infinite shine/pulse
* keyframes, and Framer Motion layout transitions all ignore the media query
* by default. Forcing every animation/transition to ~0ms collapses zoom-ins,
* slide-ins, and continuous shimmers to a single frame without removing the
* 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
* indicators across Studio (tool execution loaders, sonner toasts, Tauri
* startup / update screens, the <Spinner /> 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.
*/
@media (prefers-reduced-motion: reduce) {
*,
*::before,
*::after {
animation-duration: 0.01ms !important;
animation-iteration-count: 1 !important;
transition-duration: 0.01ms !important;
scroll-behavior: auto !important;
}
.animate-spin {
animation-duration: 1.5s !important;
animation-iteration-count: infinite !important;
}
}

57
studio/frontend/src/lib/toast.ts
View file

@ -1,59 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// sonner `toast` wrapper that defaults `dismissible: false` so swipe
// capture doesn't block text selection. Drop-in for `from "sonner"`.
import { toast as sonnerToast, type ExternalToast } from "sonner";
type AnyFn = (...args: unknown[]) => unknown;
function withDismissibleFalse<F extends AnyFn>(fn: F): F {
return ((...args: unknown[]) => {
// Branch by arity: React-element messages are objects too.
if (args.length <= 1) {
args.push({ dismissible: false } satisfies ExternalToast);
} else {
const lastIdx = args.length - 1;
const last = args[lastIdx];
if (last && typeof last === "object" && !Array.isArray(last)) {
const opts = last as ExternalToast;
if (!("dismissible" in opts)) {
args[lastIdx] = { dismissible: false, ...opts };
}
}
}
return fn(...args);
}) as F;
}
const wrappedCallable = withDismissibleFalse(
sonnerToast as unknown as AnyFn,
) as typeof sonnerToast;
// `promise(p, data?)` carries `dismissible` at the top of `data`,
// covering loading / success / error states. `dismiss`, `getHistory`,
// `getToasts` take no options.
const wrappedPromise: typeof sonnerToast.promise = ((promise, data) => {
const merged =
data && typeof data === "object" && !("dismissible" in data)
? { dismissible: false, ...data }
: (data ?? { dismissible: false });
return sonnerToast.promise(promise, merged);
}) as typeof sonnerToast.promise;
export const toast: typeof sonnerToast = Object.assign(wrappedCallable, {
success: withDismissibleFalse(sonnerToast.success.bind(sonnerToast) as AnyFn) as typeof sonnerToast.success,
info: withDismissibleFalse(sonnerToast.info.bind(sonnerToast) as AnyFn) as typeof sonnerToast.info,
warning: withDismissibleFalse(sonnerToast.warning.bind(sonnerToast) as AnyFn) as typeof sonnerToast.warning,
error: withDismissibleFalse(sonnerToast.error.bind(sonnerToast) as AnyFn) as typeof sonnerToast.error,
message: withDismissibleFalse(sonnerToast.message.bind(sonnerToast) as AnyFn) as typeof sonnerToast.message,
loading: withDismissibleFalse(sonnerToast.loading.bind(sonnerToast) as AnyFn) as typeof sonnerToast.loading,
custom: withDismissibleFalse(sonnerToast.custom.bind(sonnerToast) as AnyFn) as typeof sonnerToast.custom,
promise: wrappedPromise,
dismiss: sonnerToast.dismiss.bind(sonnerToast) as typeof sonnerToast.dismiss,
getHistory: sonnerToast.getHistory.bind(sonnerToast) as typeof sonnerToast.getHistory,
getToasts: sonnerToast.getToasts.bind(sonnerToast) as typeof sonnerToast.getToasts,
});
// Re-export of sonner. Swipe blocking lives on the Toaster via
// `swipeDirections={[]}`, so no per-toast dismissible override.
export { toast } from "sonner";
export type { ExternalToast } from "sonner";

233
studio/package-lock.json generated Normal file
View file

@ -0,0 +1,233 @@
{
"name": "unsloth-studio-tauri-cli",
"version": "0.0.0",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "unsloth-studio-tauri-cli",
"version": "0.0.0",
"license": "AGPL-3.0-only",
"devDependencies": {
"@tauri-apps/cli": "2.10.1"
}
},
"node_modules/@tauri-apps/cli": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli/-/cli-2.10.1.tgz",
"integrity": "sha512-jQNGF/5quwORdZSSLtTluyKQ+o6SMa/AUICfhf4egCGFdMHqWssApVgYSbg+jmrZoc8e1DscNvjTnXtlHLS11g==",
"dev": true,
"license": "Apache-2.0 OR MIT",
"bin": {
"tauri": "tauri.js"
},
"engines": {
"node": ">= 10"
},
"funding": {
"type": "opencollective",
"url": "https://opencollective.com/tauri"
},
"optionalDependencies": {
"@tauri-apps/cli-darwin-arm64": "2.10.1",
"@tauri-apps/cli-darwin-x64": "2.10.1",
"@tauri-apps/cli-linux-arm-gnueabihf": "2.10.1",
"@tauri-apps/cli-linux-arm64-gnu": "2.10.1",
"@tauri-apps/cli-linux-arm64-musl": "2.10.1",
"@tauri-apps/cli-linux-riscv64-gnu": "2.10.1",
"@tauri-apps/cli-linux-x64-gnu": "2.10.1",
"@tauri-apps/cli-linux-x64-musl": "2.10.1",
"@tauri-apps/cli-win32-arm64-msvc": "2.10.1",
"@tauri-apps/cli-win32-ia32-msvc": "2.10.1",
"@tauri-apps/cli-win32-x64-msvc": "2.10.1"
}
},
"node_modules/@tauri-apps/cli-darwin-arm64": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.10.1.tgz",
"integrity": "sha512-Z2OjCXiZ+fbYZy7PmP3WRnOpM9+Fy+oonKDEmUE6MwN4IGaYqgceTjwHucc/kEEYZos5GICve35f7ZiizgqEnQ==",
"cpu": [
"arm64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-darwin-x64": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-darwin-x64/-/cli-darwin-x64-2.10.1.tgz",
"integrity": "sha512-V/irQVvjPMGOTQqNj55PnQPVuH4VJP8vZCN7ajnj+ZS8Kom1tEM2hR3qbbIRoS3dBKs5mbG8yg1WC+97dq17Pw==",
"cpu": [
"x64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-linux-arm-gnueabihf": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm-gnueabihf/-/cli-linux-arm-gnueabihf-2.10.1.tgz",
"integrity": "sha512-Hyzwsb4VnCWKGfTw+wSt15Z2pLw2f0JdFBfq2vHBOBhvg7oi6uhKiF87hmbXOBXUZaGkyRDkCHsdzJcIfoJC2w==",
"cpu": [
"arm"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-linux-arm64-gnu": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-gnu/-/cli-linux-arm64-gnu-2.10.1.tgz",
"integrity": "sha512-OyOYs2t5GkBIvyWjA1+h4CZxTcdz1OZPCWAPz5DYEfB0cnWHERTnQ/SLayQzncrT0kwRoSfSz9KxenkyJoTelA==",
"cpu": [
"arm64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-linux-arm64-musl": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-arm64-musl/-/cli-linux-arm64-musl-2.10.1.tgz",
"integrity": "sha512-MIj78PDDGjkg3NqGptDOGgfXks7SYJwhiMh8SBoZS+vfdz7yP5jN18bNaLnDhsVIPARcAhE1TlsZe/8Yxo2zqg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-linux-riscv64-gnu": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-riscv64-gnu/-/cli-linux-riscv64-gnu-2.10.1.tgz",
"integrity": "sha512-X0lvOVUg8PCVaoEtEAnpxmnkwlE1gcMDTqfhbefICKDnOTJ5Est3qL0SrWxizDackIOKBcvtpejrSiVpuJI1kw==",
"cpu": [
"riscv64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-linux-x64-gnu": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.10.1.tgz",
"integrity": "sha512-2/12bEzsJS9fAKybxgicCDFxYD1WEI9kO+tlDwX5znWG2GwMBaiWcmhGlZ8fi+DMe9CXlcVarMTYc0L3REIRxw==",
"cpu": [
"x64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-linux-x64-musl": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.10.1.tgz",
"integrity": "sha512-Y8J0ZzswPz50UcGOFuXGEMrxbjwKSPgXftx5qnkuMs2rmwQB5ssvLb6tn54wDSYxe7S6vlLob9vt0VKuNOaCIQ==",
"cpu": [
"x64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-win32-arm64-msvc": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-arm64-msvc/-/cli-win32-arm64-msvc-2.10.1.tgz",
"integrity": "sha512-iSt5B86jHYAPJa/IlYw++SXtFPGnWtFJriHn7X0NFBVunF6zu9+/zOn8OgqIWSl8RgzhLGXQEEtGBdR4wzpVgg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-win32-ia32-msvc": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-ia32-msvc/-/cli-win32-ia32-msvc-2.10.1.tgz",
"integrity": "sha512-gXyxgEzsFegmnWywYU5pEBURkcFN/Oo45EAwvZrHMh+zUSEAvO5E8TXsgPADYm31d1u7OQU3O3HsYfVBf2moHw==",
"cpu": [
"ia32"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">= 10"
}
},
"node_modules/@tauri-apps/cli-win32-x64-msvc": {
"version": "2.10.1",
"resolved": "https://registry.npmjs.org/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.10.1.tgz",
"integrity": "sha512-6Cn7YpPFwzChy0ERz6djKEmUehWrYlM+xTaNzGPgZocw3BD7OfwfWHKVWxXzdjEW2KfKkHddfdxK1XXTYqBRLg==",
"cpu": [
"x64"
],
"dev": true,
"license": "Apache-2.0 OR MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">= 10"
}
}
}
}

10
studio/package.json Normal file
View file

@ -0,0 +1,10 @@
{
"name": "unsloth-studio-tauri-cli",
"version": "0.0.0",
"private": true,
"description": "Lockfile holder for @tauri-apps/cli used by the desktop release workflow. Not a real npm package; `npm ci --prefix studio` resolves the pinned Tauri CLI from this directory's package-lock.json.",
"license": "AGPL-3.0-only",
"devDependencies": {
"@tauri-apps/cli": "2.10.1"
}
}

37
studio/setup.ps1
View file

@ -1842,6 +1842,21 @@ if ($CuTag -eq "cpu") {
}
}
# Rename running unsloth.exe so pip can replace it (Windows refuses to delete a mapped .exe).
$VenvScriptsDir = Join-Path $VenvDir "Scripts"
$RunningUnslothExe = Join-Path $VenvScriptsDir "unsloth.exe"
if (Test-Path -LiteralPath $RunningUnslothExe -PathType Leaf) {
$StaleUnslothExe = "$RunningUnslothExe.deleteme"
if (Test-Path -LiteralPath $StaleUnslothExe) {
Remove-Item -LiteralPath $StaleUnslothExe -Force -ErrorAction SilentlyContinue
}
try {
Rename-Item -LiteralPath $RunningUnslothExe -NewName "unsloth.exe.deleteme" -Force -ErrorAction Stop
} catch {
substep "could not rename unsloth.exe ($($_.Exception.Message)); pip may fail with WinError 32" "Yellow"
}
}
# Ordered heavy dependency installation -- shared cross-platform script
substep "running ordered dependency installation..."
python "$PSScriptRoot\install_python_stack.py"
@ -1851,6 +1866,28 @@ $ErrorActionPreference = $prevEAP
if ($stackExit -ne 0) {
Write-Host "[FAILED] Python dependency installation failed (exit code $stackExit)" -ForegroundColor Red
Write-Host " Re-run the installer or check the error above for details." -ForegroundColor Red
# Restore the pre-rename unsloth.exe so the user keeps a working CLI.
# Treat a zero-byte exe as "pip half-wrote a broken binary" -- prefer the
# stale-but-working copy in .deleteme.
if (Test-Path -LiteralPath "$RunningUnslothExe.deleteme") {
$needRestore = -not (Test-Path -LiteralPath $RunningUnslothExe)
if (-not $needRestore) {
try {
$needRestore = (Get-Item -LiteralPath $RunningUnslothExe -ErrorAction Stop).Length -eq 0
} catch { $needRestore = $true }
}
if ($needRestore) {
try {
if (Test-Path -LiteralPath $RunningUnslothExe) {
Remove-Item -LiteralPath $RunningUnslothExe -Force -ErrorAction SilentlyContinue
}
Rename-Item -LiteralPath "$RunningUnslothExe.deleteme" -NewName "unsloth.exe" -Force -ErrorAction Stop
substep "restored unsloth.exe after failed install"
} catch {
substep "could not restore unsloth.exe ($($_.Exception.Message))" "Yellow"
}
}
}
exit 1
}

4
studio/src-tauri/src/update.rs
View file

@ -65,6 +65,10 @@ fn spawn_update(
// the same install the desktop app uses, not an inherited custom root.
cmd.env_remove("UNSLOTH_STUDIO_HOME");
cmd.env_remove("STUDIO_HOME");
// Signal to unsloth_cli that this update was initiated by the Tauri
// desktop bundle so it skips re-creating CLI launchers/.app/.desktop
// shortcuts (Tauri owns its own bundle entries).
cmd.env("UNSLOTH_TAURI_UPDATE", "1");
#[cfg(windows)]
let mut child: Box<dyn ChildWrapper + Send> = {

105
tests/security/test_lockfile_supply_chain_audit.py
View file

@ -27,9 +27,12 @@ def _run_auditor(
root: Path,
npm_lockfiles: list[Path] | None = None,
cargo_lockfiles: list[Path] | None = None,
strict: bool = False,
timeout: int = 30,
) -> subprocess.CompletedProcess:
cmd = [sys.executable, str(SCRIPT), "--root", str(root)]
if strict:
cmd.append("--strict")
for p in npm_lockfiles or []:
cmd.extend(["--npm-lockfile", str(p)])
for p in cargo_lockfiles or []:
@ -170,6 +173,32 @@ checksum = "0000000000000000000000000000000000000000000000000000000000000000"
def test_malicious_cargo_lockfile_refused(tmp_path):
"""Inline Cargo.lock with `source = "git+https://example.com/..."`
must trip the `non-registry-cargo-source` check.
`non-registry-cargo-source` is an advisory finding kind in the
auditor's default mode (per the audit script's BLOCKING_KINDS
set). To exercise the historical "refuse to install" behavior we
pass --strict here; that promotes every finding to blocking and
keeps the test honest about its intent (detection + refusal).
"""
lockfile = tmp_path / "Cargo.lock"
lockfile.write_text(_MALICIOUS_CARGO_LOCK)
proc = _run_auditor(
root = tmp_path,
npm_lockfiles = [FIXTURES / "clean_lockfile.json"],
cargo_lockfiles = [lockfile],
strict = True,
)
assert proc.returncode == 1
combined = proc.stdout + proc.stderr
assert "non-registry-cargo-source" in combined
assert "git+https://example.com" in combined
def test_malicious_cargo_lockfile_default_mode_advisory(tmp_path):
"""Default (non-strict) mode classifies `non-registry-cargo-source`
as advisory: the finding is still emitted as a `::warning::`
annotation but the process exits 0 so the build is not gated.
Regression test for the advisory/strict split.
"""
lockfile = tmp_path / "Cargo.lock"
lockfile.write_text(_MALICIOUS_CARGO_LOCK)
@ -178,10 +207,13 @@ def test_malicious_cargo_lockfile_refused(tmp_path):
npm_lockfiles = [FIXTURES / "clean_lockfile.json"],
cargo_lockfiles = [lockfile],
)
assert proc.returncode == 1
assert proc.returncode == 0, (
f"expected exit 0 (advisory), got {proc.returncode}\n"
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}"
)
combined = proc.stdout + proc.stderr
assert "non-registry-cargo-source" in combined
assert "git+https://example.com" in combined
assert "advisory finding" in combined
def test_audit_cargo_lockfile_direct_call(tmp_path):
@ -192,6 +224,75 @@ def test_audit_cargo_lockfile_direct_call(tmp_path):
assert "non-registry-cargo-source" in kinds
# ---------------------------------------------------------------------------
# GitHub Actions annotation escape: ::warning:: / ::error:: messages
# are truncated at the first newline unless escaped, so the multi-line
# Finding must be collapsed via the spec'd %0A / %0D / %25 encoding.
# ---------------------------------------------------------------------------
def test_gha_escape_collapses_finding_to_one_line():
"""`_gha_escape()` must collapse newlines (`%0A`), carriage
returns (`%0D`), and percent signs (`%25`) so that
`::warning::<msg>` / `::error::<msg>` render the full finding
in the GitHub Actions UI annotation instead of being truncated
at the first newline. The `%` replacement must happen first or
the subsequent `%0A` / `%0D` escapes get double-encoded.
"""
assert lsa._gha_escape("a\nb\nc") == "a%0Ab%0Ac"
assert lsa._gha_escape("a\rb") == "a%0Db"
assert lsa._gha_escape("100%") == "100%25"
# Order regression: `%` must escape before `\n` so the literal
# text `a%b\nc` becomes `a%25b%0Ac`, not `a%250Ab%0Ac`.
assert lsa._gha_escape("a%b\nc") == "a%25b%0Ac"
f = lsa.Finding(
path = "/x/lock.json",
package = "node_modules/foo",
kind = "missing-integrity-hash",
detail = "bad stuff",
)
escaped = lsa._gha_escape(str(f))
assert "\n" not in escaped
assert "%0A" in escaped
assert "missing-integrity-hash" in escaped
assert "node_modules/foo" in escaped
assert "bad stuff" in escaped
def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
"""End-to-end check: the `::warning::` line emitted for an
advisory finding must be a SINGLE physical line (the rest of
the Finding is `%0A`-escaped inside the message). Regression
test for the gemini-code-assist review on PR #5604: without
`_gha_escape`, GitHub Actions truncates the annotation after
`[kind] path` and the package + detail fields never render.
"""
lockfile = tmp_path / "Cargo.lock"
lockfile.write_text(_MALICIOUS_CARGO_LOCK)
proc = _run_auditor(
root = tmp_path,
npm_lockfiles = [FIXTURES / "clean_lockfile.json"],
cargo_lockfiles = [lockfile],
)
warning_lines = [
line for line in proc.stderr.splitlines() if line.startswith("::warning::")
]
assert warning_lines, (
"expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}"
)
for line in warning_lines:
# Single physical line: kind, package, detail all present
# via %0A escape, not split across stderr lines.
assert "%0A" in line, (
f"::warning:: line has no %0A escape; multi-line text "
f"would be truncated by GH Actions:\n{line}"
)
assert "non-registry-cargo-source" in line
assert "package:" in line
assert "detail:" in line
# ---------------------------------------------------------------------------
# SF4: skip env var requires a justification value.
# ---------------------------------------------------------------------------

48
tests/studio/playwright_chat_ui.py
View file

@ -113,6 +113,49 @@ def fail(m):
raise AssertionError(f"[ui] FAIL: {m}")
def expected_default_model():
override = os.environ.get("EXPECTED_DEFAULT_MODEL")
if override:
return override
# Parse DEFAULT_MODELS_GGUF as a literal out of defaults.py instead of
# importing it. The Playwright job installs Studio with --no-torch, so
# the studio.backend.core.inference package init (which eagerly imports
# the orchestrator -> structlog) and defaults.py's own
# `import utils.hardware.hardware as hw` are both unavailable.
import ast
defaults_path = (
Path(__file__).resolve().parents[2]
/ "studio"
/ "backend"
/ "core"
/ "inference"
/ "defaults.py"
)
try:
tree = ast.parse(defaults_path.read_text())
except Exception as exc:
fail(f"could not read {defaults_path}: {exc}")
models = None
for node in tree.body:
if not isinstance(node, ast.Assign):
continue
if not any(
isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF"
for t in node.targets
):
continue
try:
models = ast.literal_eval(node.value)
except Exception as exc:
fail(f"could not eval DEFAULT_MODELS_GGUF literal: {exc}")
break
if not models:
fail("DEFAULT_MODELS_GGUF not found or empty in defaults.py")
return models[0]
def soft_fail(m):
"""Hard fail in STRICT mode, info-warn otherwise.
@ -475,10 +518,7 @@ with sync_playwright() as p:
# list or hides the default would break the first-launch UX,
# which is what this assertion guards.
step("default_models[0] matches DEFAULT_MODELS_GGUF[0]")
EXPECTED_DEFAULT = os.environ.get(
"EXPECTED_DEFAULT_MODEL",
"unsloth/gemma-4-E2B-it-GGUF",
)
EXPECTED_DEFAULT = expected_default_model()
defaults_resp = evaluate_fetch(
page,
f"{BASE}/api/models/list",

92
tests/test_finetune_last_n_layers.py Normal file
View file

@ -0,0 +1,92 @@
# Unsloth - 2x faster, 70% less memory LLM finetuning
# Tests for the `finetune_last_n_layers` parity knob (CUDA side).
#
# Mirrors unsloth-zoo's `FastMLXModel.get_peft_model` parameter.
# mlx-lm CLI's CONFIG_DEFAULTS['num_layers']=16 applies LoRA to the
# last 16 transformer blocks only. On the CUDA path, PEFT exposes
# `layers_to_transform` to do the same. This convenience knob fills
# `layers_to_transform` for the user when set, matching mlx-lm CLI
# AND unsloth-zoo's MLX path with a single config value.
#
# The tests intentionally avoid pulling in CUDA / a real model
# checkpoint — they exercise only the helper that translates
# `finetune_last_n_layers` into `layers_to_transform`.
from __future__ import annotations
import pytest
def test_get_total_transformer_layers_reads_num_hidden_layers():
from unsloth.models.vision import _get_total_transformer_layers
class FakeConfig:
num_hidden_layers = 18
class FakeModel:
config = FakeConfig()
assert _get_total_transformer_layers(FakeModel()) == 18
def test_get_total_transformer_layers_reads_text_config():
from unsloth.models.vision import _get_total_transformer_layers
class TextConfig:
num_hidden_layers = 24
class FakeConfig:
text_config = TextConfig()
class FakeModel:
config = FakeConfig()
# No num_hidden_layers at top level — should fall through to text_config.
assert _get_total_transformer_layers(FakeModel()) == 24
def test_get_total_transformer_layers_handles_alternative_attr_names():
from unsloth.models.vision import _get_total_transformer_layers
for attr in ("n_layer", "n_layers", "num_layers"):
cfg = type("Cfg", (), {attr: 12})()
model = type("M", (), {"config": cfg})()
assert _get_total_transformer_layers(model) == 12
def test_get_total_transformer_layers_returns_none_when_unknown():
from unsloth.models.vision import _get_total_transformer_layers
class FakeConfig:
pass
class FakeModel:
config = FakeConfig()
assert _get_total_transformer_layers(FakeModel()) is None
def test_get_total_transformer_layers_returns_none_for_missing_config():
from unsloth.models.vision import _get_total_transformer_layers
class FakeModel:
pass
assert _get_total_transformer_layers(FakeModel()) is None
def test_finetune_last_n_layers_signature_present_on_llama_and_vision():
"""Both entry points must expose the new parameter with default None."""
import inspect
from unsloth.models.llama import FastLlamaModel
from unsloth.models.vision import FastBaseModel
for cls in (FastLlamaModel, FastBaseModel):
sig = inspect.signature(cls.get_peft_model)
assert (
"finetune_last_n_layers" in sig.parameters
), f"{cls.__name__}.get_peft_model missing finetune_last_n_layers"
assert sig.parameters["finetune_last_n_layers"].default is None, (
f"{cls.__name__}.get_peft_model: finetune_last_n_layers default "
f"must be None to preserve historical behavior"
)

67
tests/test_import_fixes_drift.py
View file

@ -545,6 +545,73 @@ def test_transformers_pretrained_model_has_get_input_embeddings():
# ===========================================================================
# ===========================================================================
# transformers LOSS_MAPPING -- patch_loss_functions() coverage
# Regression for https://github.com/unslothai/unsloth/issues/4188:
# Qwen3_5ForConditionalGeneration has loss_type='ForConditionalGeneration',
# a separate LOSS_MAPPING key that was never patched, leaving the model with
# the stock ForCausalLMLoss which does logits.float() and OOMs on <=24 GB GPUs.
# ===========================================================================
def _reset_loss_mapping(mapping, saved):
mapping.clear()
mapping.update(saved)
def test_patch_loss_functions_covers_conditional_generation():
"""After patch_loss_functions(), every LOSS_MAPPING key that was aliased
to ForCausalLMLoss must also point at the Unsloth kernel -- not just
LOSS_MAPPING['ForCausalLM']."""
lu = pytest.importorskip("transformers.loss.loss_utils")
cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")
saved = dict(lu.LOSS_MAPPING)
try:
cel.patch_loss_functions(torch_compile = False)
unsloth_loss = lu.LOSS_MAPPING.get("ForCausalLM")
assert unsloth_loss is not None
assert "Unsloth" in str(
unsloth_loss
), f"LOSS_MAPPING['ForCausalLM'] was not replaced: {unsloth_loss}"
cg_loss = lu.LOSS_MAPPING.get("ForConditionalGeneration")
assert cg_loss is unsloth_loss, (
f"LOSS_MAPPING['ForConditionalGeneration'] not patched: {cg_loss}. "
f"Qwen3_5ForConditionalGeneration will silently use the stock "
f"ForCausalLMLoss and OOM at large sequence lengths."
)
finally:
_reset_loss_mapping(lu.LOSS_MAPPING, saved)
def test_patch_loss_functions_does_not_touch_other_loss_types():
"""patch_loss_functions() must not overwrite unrelated loss types
(segmentation, detection, masked-LM, etc.) with the causal-LM kernel."""
lu = pytest.importorskip("transformers.loss.loss_utils")
cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")
non_causal_keys = {
k
for k, v in lu.LOSS_MAPPING.items()
if getattr(v, "__name__", "") != "ForCausalLMLoss"
}
saved = dict(lu.LOSS_MAPPING)
try:
cel.patch_loss_functions(torch_compile = False)
unsloth_loss = lu.LOSS_MAPPING.get("ForCausalLM")
for key in non_causal_keys:
assert lu.LOSS_MAPPING.get(key) is not unsloth_loss, (
f"patch_loss_functions() incorrectly overwrote "
f"LOSS_MAPPING['{key}'] with the Unsloth ForCausalLM kernel."
)
finally:
_reset_loss_mapping(lu.LOSS_MAPPING, saved)
def test_accelerate_utils_imports_module_present():
"""``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` (import_fixes.py
493-516, 1320-1372). Both reach into accelerate.utils.imports."""

7
unsloth/dataprep/synthetic.py
View file

@ -21,7 +21,12 @@ from collections import deque
import time
import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
_OFFLINE_VALS = {"1", "true", "yes", "on"}
if not (
os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_VALS
or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_VALS
):
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
import requests
import torch
import gc

15
unsloth/kernels/cross_entropy_loss.py
View file

@ -461,3 +461,18 @@ if (Version(torch.__version__) < Version("2.4.0")) and not hasattr(
# Patch CE Losses in transformers
def patch_loss_functions(torch_compile = True):
_patch_loss_functions(fast_cross_entropy_loss, torch_compile = torch_compile)
# Defense-in-depth sweep for LOSS_MAPPING aliases still pointing at the
# stock ForCausalLMLoss (e.g. ForConditionalGeneration for Qwen3.5,
# CsmForConditionalGeneration). unsloth_zoo also does this; remove once
# the floor pin moves past unslothai/unsloth-zoo#656.
try:
import transformers.loss.loss_utils as _lu
_unsloth_loss = _lu.LOSS_MAPPING.get("ForCausalLM")
if _unsloth_loss is not None:
for _key, _fn in list(_lu.LOSS_MAPPING.items()):
if getattr(_fn, "__name__", "") == "ForCausalLMLoss":
_lu.LOSS_MAPPING[_key] = _unsloth_loss
except (ImportError, AttributeError):
pass

9
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2026.5.4"
__version__ = "2026.5.5"
__all__ = [
"SUPPORTS_BFLOAT16",
@ -1771,6 +1771,13 @@ def get_statistics(local_files_only = False):
return
if local_files_only:
return
# Also skip when HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE are set.
_offline_vals = {"1", "true", "yes", "on"}
if (
os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline_vals
or os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline_vals
):
return
from huggingface_hub.utils import (
disable_progress_bars,
enable_progress_bars,

10
unsloth/models/llama.py
View file

@ -2831,6 +2831,7 @@ class FastLlamaModel:
bias = "none",
layers_to_transform = None,
layers_pattern = None,
finetune_last_n_layers = None,
use_gradient_checkpointing = "unsloth",
random_state = 3407,
max_seq_length = 2048, # not used anymore
@ -2863,6 +2864,7 @@ class FastLlamaModel:
bias = bias,
layers_to_transform = layers_to_transform,
layers_pattern = layers_pattern,
finetune_last_n_layers = finetune_last_n_layers,
use_gradient_checkpointing = use_gradient_checkpointing,
random_state = random_state,
max_seq_length = max_seq_length,
@ -3160,6 +3162,14 @@ class FastLlamaModel:
if target_parameters is None:
target_parameters = get_moe_target_parameters(model, target_modules)
if finetune_last_n_layers is not None and layers_to_transform is None:
from .vision import _get_total_transformer_layers
_total_layers = _get_total_transformer_layers(model)
if _total_layers is not None and _total_layers > 0:
_n = max(1, min(int(finetune_last_n_layers), _total_layers))
layers_to_transform = list(range(_total_layers - _n, _total_layers))
arguments = dict(
r = r,
lora_alpha = lora_alpha,

45
unsloth/models/loader.py
View file

@ -99,15 +99,22 @@ from ._utils import (
fast_inference_setup,
)
global FORCE_FLOAT32
# Forces float32 precision since float16 goes to infinity
FORCE_FLOAT32 = [
"gemma3,", # Add comma bc gemma3 will match gemma3n
"gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3)
"gemma3n",
"gpt_oss",
"qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training
]
# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers
# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working.
# Fallback list mirrors zoo for users who upgrade unsloth without upgrading
# unsloth_zoo (so this module never fails at import).
try:
from unsloth_zoo import FORCE_FLOAT32 # noqa: F401
except ImportError:
global FORCE_FLOAT32
# Forces float32 precision since float16 goes to infinity
FORCE_FLOAT32 = [
"gemma3,", # Add comma bc gemma3 will match gemma3n
"gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3)
"gemma3n",
"gpt_oss",
"qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training
]
global DISABLE_COMPILE_MODEL_NAMES
# Must be alphabetically sorted for each entry
@ -308,6 +315,16 @@ class FastLanguageModel(FastLlamaModel):
if is_dist:
device_map = distributed_device_map
# Honour offline env vars BEFORE FastModel delegation so 8bit /
# full-finetuning / qat paths also receive local_files_only.
if not kwargs.get("local_files_only", False):
_offline = {"1", "true", "yes", "on"}
if (
os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline
or os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline
):
kwargs["local_files_only"] = True
if load_in_8bit or full_finetuning or qat_scheme is not None:
return FastModel.from_pretrained(
model_name = model_name,
@ -1055,6 +1072,15 @@ class FastModel(FastBaseModel):
model_config = None
peft_config = None
local_files_only = kwargs.get("local_files_only", False)
# Mirror env-var fallback for direct callers (FastVisionModel / FastTextModel).
if not local_files_only:
_offline = {"1", "true", "yes", "on"}
if (
os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline
or os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline
):
local_files_only = True
kwargs["local_files_only"] = True
try:
model_config = AutoConfig.from_pretrained(
@ -1362,7 +1388,6 @@ class FastModel(FastBaseModel):
for model_type_arch in model_types:
if model_type_arch != "siglip":
break
global FORCE_FLOAT32
for disable_name in FORCE_FLOAT32:
# add comma to model_types_all matching in case of exact match for end
if (

36
unsloth/models/vision.py
View file

@ -547,6 +547,35 @@ def _construct_vlm_processor_fallback(
return None
def _get_total_transformer_layers(model):
"""Best-effort total transformer block count across HF model shapes.
Returns None if not determinable; caller should skip the conversion."""
cfg = getattr(model, "config", None)
if cfg is None:
return None
for name in (
"num_hidden_layers",
"n_layer",
"n_layers",
"num_layers",
):
v = getattr(cfg, name, None)
if isinstance(v, int) and v > 0:
return v
text_cfg = getattr(cfg, "text_config", None)
if text_cfg is not None:
for name in (
"num_hidden_layers",
"n_layer",
"n_layers",
"num_layers",
):
v = getattr(text_cfg, name, None)
if isinstance(v, int) and v > 0:
return v
return None
class FastBaseModel:
@staticmethod
def from_pretrained(
@ -1319,6 +1348,7 @@ class FastBaseModel:
finetune_language_layers = True,
finetune_attention_modules = True,
finetune_mlp_modules = True,
finetune_last_n_layers = None,
layers_to_transform = None,
layers_pattern = None,
use_gradient_checkpointing = "unsloth",
@ -1417,6 +1447,12 @@ class FastBaseModel:
if target_parameters is None:
target_parameters = get_moe_target_parameters(model, target_modules)
if finetune_last_n_layers is not None and layers_to_transform is None:
_total_layers = _get_total_transformer_layers(model)
if _total_layers is not None and _total_layers > 0:
n = max(1, min(int(finetune_last_n_layers), _total_layers))
layers_to_transform = list(range(_total_layers - n, _total_layers))
# Get only allowed parameters for LoraConfig
local_variables = {
**locals(),

253
unsloth_cli/commands/studio.py
View file

@ -6,6 +6,7 @@ import hashlib
import json
import os
import platform
import re
import secrets
import sqlite3
import subprocess
@ -13,6 +14,8 @@ import sys
import tempfile
import time
import types
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional
@ -1064,6 +1067,170 @@ def _run_setup_script(*, verbose: bool = False) -> None:
raise typer.Exit(result.returncode)
_INSTALLER_URL_BASH = "https://unsloth.ai/install.sh"
_INSTALLER_URL_PWSH = "https://unsloth.ai/install.ps1"
def _refresh_desktop_shortcuts(*, verbose: bool = False) -> None:
"""Re-run installer with --shortcuts-only to refresh launchers post-update."""
env = {**os.environ}
if verbose:
env["UNSLOTH_VERBOSE"] = "1"
is_windows = platform.system() == "Windows"
installer_name = "install.ps1" if is_windows else "install.sh"
installer_url = _INSTALLER_URL_PWSH if is_windows else _INSTALLER_URL_BASH
# Prefer local checkout, fall back to package dir, then network fetch.
local_repo = (os.environ.get("STUDIO_LOCAL_REPO") or "").strip()
candidates: list[Path] = []
if local_repo:
candidates.append(Path(local_repo) / installer_name)
candidates.append(_PACKAGE_ROOT / installer_name)
args = ["--shortcuts-only"]
if verbose:
args.append("--verbose")
if is_windows:
ps_argv: list[str] = ["powershell.exe"]
if _should_hide_windows_subprocesses():
ps_argv.extend(
["-NoLogo", "-NoProfile", "-NonInteractive", "-WindowStyle", "Hidden"]
)
for script in candidates:
try:
if script.is_file():
quoted = str(script).replace("'", "''")
argv = list(ps_argv)
argv.extend(
[
"-ExecutionPolicy",
"Bypass",
"-Command",
f"& '{quoted}' {' '.join(args)} *>&1",
]
)
result = subprocess.run(
argv,
env = env,
check = False,
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
typer.echo(
f" refresh-launcher install.ps1 exited {result.returncode}"
)
return
except OSError:
continue
# PyPI installs lack install.ps1: fetch + pipe to powershell stdin.
try:
request = urllib.request.Request(
installer_url, headers = {"User-Agent": "unsloth-studio-update"}
)
with urllib.request.urlopen(request, timeout = 30) as response:
installer = response.read().decode("utf-8", errors = "replace")
except (urllib.error.URLError, TimeoutError, OSError) as exc:
typer.echo(
f" refresh-launcher skipped: could not fetch {installer_url} ({exc})"
)
return
# install.ps1 auto-invokes `Install-UnslothStudio @args` at EOF; over
# stdin `$args` is empty so that triggers the full installer flow
# (deps, venv, prompts) before our shortcuts-only call. Strip it.
installer = re.sub(
r"(?m)^[ \t]*Install-UnslothStudio[ \t]+@args[ \t]*\r?\n?",
"",
installer,
)
# stdin-piped scripts have empty $args, so call Install-UnslothStudio explicitly.
marker_args = " ".join(args)
wrapper = installer + f"\nInstall-UnslothStudio {marker_args}\n"
# Write to a UTF-8 BOM tempfile and use -File rather than -Command -.
# `powershell.exe -Command -` reads stdin via [Console]::InputEncoding
# (CP1252/OEM on most Windows boxes), which mangles box-drawing chars
# in install.ps1. -File reads the BOM and decodes correctly. The
# prefix gives AV/EDR engines (and grep'ing users) a clear identity.
ps1_fd, ps1_path = tempfile.mkstemp(
prefix = "unsloth-studio-refresh-",
suffix = ".ps1",
)
try:
with os.fdopen(ps1_fd, "wb") as fh:
fh.write(b"\xef\xbb\xbf" + wrapper.encode("utf-8"))
argv = list(ps_argv)
argv.extend(["-ExecutionPolicy", "Bypass", "-File", ps1_path])
try:
result = subprocess.run(
argv,
env = env,
check = False,
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
typer.echo(
f" refresh-launcher fetched install.ps1 exited {result.returncode}"
)
except OSError as exc:
typer.echo(
f" refresh-launcher skipped: powershell exec failed ({exc})"
)
finally:
try:
os.unlink(ps1_path)
except OSError:
pass
return
for script in candidates:
try:
if script.is_file():
result = subprocess.run(
["bash", str(script), *args],
env = env,
check = False,
)
if result.returncode != 0:
typer.echo(
f" refresh-launcher install.sh exited {result.returncode}"
)
return
except OSError:
continue
# PyPI installs lack install.sh: fetch upstream.
try:
request = urllib.request.Request(
installer_url, headers = {"User-Agent": "unsloth-studio-update"}
)
with urllib.request.urlopen(request, timeout = 30) as response:
installer = response.read()
except (urllib.error.URLError, TimeoutError, OSError) as exc:
typer.echo(
f" refresh-launcher skipped: could not fetch {installer_url} ({exc})"
)
return
try:
result = subprocess.run(
["bash", "-s", "--", *args],
input = installer,
env = env,
check = False,
)
if result.returncode != 0:
typer.echo(
f" refresh-launcher fetched install.sh exited {result.returncode}"
)
except OSError as exc:
typer.echo(f" refresh-launcher skipped: bash exec failed ({exc})")
@studio_app.command(hidden = True)
def setup(
verbose: bool = typer.Option(
@ -1093,6 +1260,9 @@ def update(
),
):
"""Update Unsloth Studio dependencies and rebuild."""
# Re-export UNSLOTH_STUDIO_HOME for env-mode installs so the refresh
# subprocess resolves the same install root the user originally chose.
_ensure_studio_env_exported()
# Ensure SKIP_STUDIO_BASE is not inherited from a parent install.ps1 session
os.environ.pop("SKIP_STUDIO_BASE", None)
os.environ["STUDIO_PACKAGE_NAME"] = package
@ -1105,7 +1275,88 @@ def update(
else:
os.environ["STUDIO_LOCAL_INSTALL"] = "0"
os.environ.pop("STUDIO_LOCAL_REPO", None)
_run_setup_script(verbose = verbose)
_release_self_exe_lock_windows()
try:
_run_setup_script(verbose = verbose)
except BaseException:
# Restore unsloth.exe from .deleteme if setup failed before pip
# produced a replacement; otherwise the user has no CLI for recovery.
_restore_self_exe_lock_windows()
raise
# On Windows clear the .deleteme orphan now that pip wrote a fresh
# unsloth.exe; on next update os.replace would overwrite it anyway,
# but leaving a stale binary around invites cross-version restore
# confusion from _restore_self_exe_lock_windows.
_cleanup_self_exe_lock_windows()
# Tauri desktop owns its own bundle entries; skip CLI launcher refresh
# so a Tauri-initiated update doesn't create duplicate shortcuts.
if os.environ.get("UNSLOTH_TAURI_UPDATE") == "1":
if verbose:
typer.echo(" refresh-launcher skipped (Tauri update)")
return
_refresh_desktop_shortcuts(verbose = verbose)
def _release_self_exe_lock_windows() -> None:
"""Rename running unsloth.exe so pip can replace it. setup.ps1 also retries."""
if platform.system() != "Windows":
return
try:
venv_scripts = Path(sys.executable).resolve().parent
except OSError:
return
exe = venv_scripts / "unsloth.exe"
if not exe.exists():
return
stale = exe.with_suffix(".exe.deleteme")
try:
# os.replace is atomic-overwrite on Windows; os.rename would raise
# FileExistsError if a prior aborted update left a .deleteme behind.
os.replace(exe, stale)
except OSError as e:
# Not fatal; setup.ps1 retries from a sibling process.
print(f"[update] could not rename {exe.name} -> {stale.name}: {e}")
def _restore_self_exe_lock_windows() -> None:
"""If setup failed before pip wrote a working unsloth.exe, restore .deleteme."""
if platform.system() != "Windows":
return
try:
venv_scripts = Path(sys.executable).resolve().parent
except OSError:
return
exe = venv_scripts / "unsloth.exe"
stale = exe.with_suffix(".exe.deleteme")
if not stale.exists():
return
# Treat a missing or zero-byte exe as "pip didn't produce a usable
# replacement"; otherwise leave the new binary alone.
if exe.exists():
try:
if exe.stat().st_size > 0:
return
except OSError:
return
try:
os.replace(stale, exe)
except OSError as e:
print(f"[update] could not restore {stale.name} -> {exe.name}: {e}")
def _cleanup_self_exe_lock_windows() -> None:
"""Remove the .deleteme orphan after a successful update on Windows."""
if platform.system() != "Windows":
return
try:
venv_scripts = Path(sys.executable).resolve().parent
except OSError:
return
stale = (venv_scripts / "unsloth.exe").with_suffix(".exe.deleteme")
try:
stale.unlink(missing_ok = True)
except OSError:
pass
# ── unsloth studio reset-password ────────────────────────────────────