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

Merge remote-tracking branch 'origin/main' into woa-nvidia-wsl-fallback

Browse source
This commit is contained in:
Daniel Han 2026-06-14 21:29:43 -07:00
commit b0ee45c34d
288 changed files with 24004 additions and 8020 deletions
Download patch file Download diff file Expand all files Collapse all files

8
.git-blame-ignore-revs Normal file
View file

@ -0,0 +1,8 @@
# Commits listed here are skipped by `git blame` so that bulk, whitespace-only
# changes don't obscure the real authorship of a line.
#
# GitHub honors this file automatically. To use it locally, run once:
# git config blame.ignoreRevsFile .git-blame-ignore-revs
# chore(studio/frontend): normalize line endings to LF
c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a

6
.gitattributes vendored
View file

@ -5,3 +5,9 @@
# clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks
# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -").
*.sh text eol=lf
# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather
# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files
# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts)
# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF.
studio/frontend/** text=auto eol=lf

61
.github/workflows/cross-platform-parity-ci.yml vendored Normal file
View file

@ -0,0 +1,61 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Runs tests/python/test_cross_platform_parity.py on Windows and macOS.
#
# Why: that test is the guard that install.sh and install.ps1 stay in
# sync, but today it only runs on ubuntu-latest (auto-discovered by
# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both
# installer scripts, and on Windows Path.read_text() defaults to the
# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already
# contains a U+274C) raises UnicodeDecodeError there even though Linux and
# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
# #6166; this job keeps that from silently regressing by exercising the
# test on the platforms it claims parity for. Pure pytest, no GPU,
# sub-second, so the matrix is cheap.
name: Cross-platform parity
on:
pull_request:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
branches: [main]
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
parity:
name: parity (${{ matrix.os }})
strategy:
fail-fast: false
matrix:
os: [windows-latest, macos-latest]
runs-on: ${{ matrix.os }}
timeout-minutes: 10
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- run: python -m pip install -U pip pytest
- name: Cross-platform parity test
run: python -m pytest tests/python/test_cross_platform_parity.py -q

6
CONTRIBUTING.md
View file

@ -27,3 +27,9 @@ Your support extends beyond code:
Finally, please be mindful of our [Code of Conduct](https://github.com/unslothai/unsloth/blob/main/CODE_OF_CONDUCT.md) to ensure a welcoming and inclusive environment for everyone.
Thank you so much for reading and we hope you have lots of fun using Unsloth! 🦥
## Pull Request Guidelines
- Keep PRs focused on a single change
- Include a concise description and motivation
- Link related issues when applicable

83
install.ps1
View file

@ -811,9 +811,11 @@ exit 0
# even when install.ps1 is executed from PowerShell 7.
$utf8Bom = New-Object System.Text.UTF8Encoding($true)
[System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom)
# shell.Run(cmd, 0, ...) already hides the window, so -WindowStyle Hidden
# is redundant; omitting it trims an AV-heuristic token (Kaspersky FP).
$vbsContent = @"
Set shell = CreateObject("WScript.Shell")
cmd = "powershell -NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File ""$launcherPs1"""
cmd = "powershell -NoProfile -ExecutionPolicy Bypass -File ""$launcherPs1"""
shell.Run cmd, 0, False
"@
# WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths.
@ -1189,14 +1191,38 @@ shell.Run cmd, 0, False
if ($SkipTorch) { $InitialGpuBranch = "no_torch" }
Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion
# ── Install uv if not present ──
# ── Install uv ──
Write-TauriLog "STEP" "Installing uv package manager"
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
substep "installing uv package manager..."
$UvMinVersion = "0.7.22"
function Test-UvVersionOk {
$cmd = Get-Command uv -ErrorAction SilentlyContinue
if (-not $cmd) { return $false }
try {
$raw = (& uv --version 2>$null | Select-Object -First 1)
} catch {
return $false
}
if ($raw -notmatch 'uv\s+([0-9]+(?:\.[0-9]+)+)') { return $false }
try {
return ([version]$Matches[1] -ge [version]$UvMinVersion)
} catch {
return $false
}
}
if (-not (Test-UvVersionOk)) {
if (Get-Command uv -ErrorAction SilentlyContinue) {
substep "updating uv package manager..."
} else {
substep "installing uv package manager..."
}
if ($script:WingetAvailable) {
$prevEAP = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {}
try { winget upgrade --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {}
if (-not (Test-UvVersionOk)) {
try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {}
}
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
}
@ -1204,19 +1230,40 @@ shell.Run cmd, 0, False
# use Astral's official PowerShell installer. This is the only
# supported path on hosts without winget (Windows ARM64 runners,
# corporate machines without the Store, etc.).
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
if (-not (Test-UvVersionOk)) {
substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow"
Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1")
Refresh-SessionPath
}
}
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
# A freshly installed uv can sit later on PATH than an older one (active
# venv, Scoop/pipx shim). Prefer a just-installed uv from a known location.
if (-not (Test-UvVersionOk)) {
$origPath = $env:PATH
foreach ($d in @($env:UV_INSTALL_DIR, $env:XDG_BIN_HOME,
(Join-Path $env:USERPROFILE ".local\bin"),
(Join-Path $env:LOCALAPPDATA "Microsoft\WinGet\Links"))) {
if ($d -and (Test-Path $d)) {
$env:PATH = "$d;$origPath"
if (Test-UvVersionOk) { break }
$env:PATH = $origPath
}
}
}
if (-not (Test-UvVersionOk)) {
step "uv" "could not be installed" "Red"
substep "Install it from https://docs.astral.sh/uv/" "Yellow"
return (Exit-InstallFailure "uv could not be installed")
}
# When bytecode compilation is enabled, large installs can exceed uv's 60s
# default on slow machines. Default to 180s, preserving overrides ("0" disables).
if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT) {
$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"
}
# ── Create venv (migrate old layout if possible, otherwise fresh) ──
# Pass the resolved executable path to uv so it does not re-resolve
# a version string back to a conda interpreter.
@ -1541,7 +1588,7 @@ shell.Run cmd, 0, False
# popping a UAC/DiskPart prompt RunAsInvoker can't suppress (manifest is
# asInvoker). So only probe when a HIP SDK is present (hipinfo found ->
# un-elevated) or the user opts in; else fall through to WMI name inference
# (enough to pick ROCm wheels + lemonade llama.cpp).
# (enough to pick ROCm wheels + the ROCm llama.cpp prebuilt).
# An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the
# HIP-SDK heuristic: a HIP SDK binary with a broken runtime can still pop the
# prompt, so $HipSdkInstalled must NOT silently re-enable it.
@ -1592,7 +1639,7 @@ shell.Run cmd, 0, False
# ── Arch resolution: env-var override → name inference ──────────────
# Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime
# ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the
# studio setup forward --rocm-gfx and pull a GPU-accelerated (lemonade)
# studio setup forward --rocm-gfx and pull a GPU-accelerated ROCm
# llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels
# still require a confirmed HIP SDK -- they stay gated on $HasROCm below.
if (-not $ROCmGfxArch) {
@ -1603,7 +1650,7 @@ shell.Run cmd, 0, False
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan"
}
# 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI).
# Targets only arches the lemonade-sdk ROCm prebuilts cover
# Targets only arches the ROCm prebuilts cover
# (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU.
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@ -1614,9 +1661,9 @@ shell.Run cmd, 0, False
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
@{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point)
@{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- lemonade gfx103X
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X
@{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X
@{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family
@{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- gfx103X family
)
foreach ($row in $nameArchTable) {
if ($ROCmGpuLabel -match $row.P) {
@ -2221,7 +2268,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.6.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2235,7 +2282,7 @@ shell.Run cmd, 0, False
}
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2282,7 +2329,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.6.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
@ -2294,7 +2341,7 @@ shell.Run cmd, 0, False
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2322,7 +2369,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.6.3" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)

84
install.sh
View file

@ -1456,7 +1456,11 @@ fi
# ── Install uv ──
tauri_log "STEP" "Installing uv package manager"
UV_MIN_VERSION="0.7.14"
UV_MIN_VERSION="0.7.22"
# When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables).
: "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"
export UV_COMPILE_BYTECODE_TIMEOUT
version_ge() {
# returns 0 if $1 >= $2
@ -2043,6 +2047,37 @@ _pick_radeon_wheel() {
# CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts
# the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 +
# librocdxg), then sources the env it persisted so detection finds the GPU.
# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d
# so non-login Studio/llama launches inherit it. Idempotent (writes only when
# the drop-in is missing); no-op without librocdxg, so never fires off WSL.
# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes
# after this shell on a non-root reinstall. Best-effort either way.
_persist_rocm_wsl_dropin() {
[ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ] || return 0
_rw_rocm=/opt/rocm
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
case ":${PATH}:" in
*":${_rw_rocm}/bin:"*) ;;
*) export PATH="${_rw_rocm}/bin:${PATH}" ;;
esac
export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}"
[ -r /etc/profile.d/unsloth-rocm-wsl.sh ] && return 0
_rw_dropin="$(
printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n'
printf 'export HSA_ENABLE_DXG_DETECTION=1\n'
printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n'
printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}"
printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}"
printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n'
)"
if [ "$(id -u)" = "0" ]; then
printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
elif command -v sudo >/dev/null 2>&1; then
printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true
fi
}
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
@ -2056,6 +2091,11 @@ _maybe_bootstrap_rocm_wsl() {
_ensure_rocm_probe_env
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
# rocminfo may work only via the transient env _ensure_rocm_probe_env
# just set, which dies with the installer. Persist the drop-in so login
# shells (Studio, llama.cpp) inherit it -- else a reinstall over an
# existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU.
_persist_rocm_wsl_dropin
return 0
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
@ -2073,31 +2113,8 @@ _maybe_bootstrap_rocm_wsl() {
. /etc/profile.d/unsloth-rocm-wsl.sh || true
else
# librocdxg present but the env drop-in is gone (e.g. a Studio
# uninstall removed it while keeping shared ROCm). Restore the FULL
# env inline (so rocminfo is on PATH) and recreate the drop-in.
_rw_rocm=/opt/rocm
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${_rw_rocm}/bin:${PATH}"
export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}"
# Persist the drop-in so later non-login Studio launches get the env
# too. /etc/profile.d is root-owned: a plain redirect fails for a
# non-root reinstall (ROCm would silently disappear after this shell),
# so tee through sudo when not root. Best-effort -- the current shell
# already has the env, so the install proceeds either way.
_rw_dropin="$(
printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n'
printf 'export HSA_ENABLE_DXG_DETECTION=1\n'
printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n'
printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}"
printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}"
printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n'
)"
if [ "$(id -u)" = "0" ]; then
printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
elif command -v sudo >/dev/null 2>&1; then
printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true
fi
# uninstall removed it while keeping shared ROCm). Restore the env.
_persist_rocm_wsl_dropin
fi
return 0
fi
@ -2405,7 +2422,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.6.3" unsloth-zoo
"unsloth>=2026.6.7" unsloth-zoo
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2418,7 +2435,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.6.3" unsloth-zoo
"unsloth>=2026.6.7" unsloth-zoo
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2622,7 +2639,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.6.3" unsloth-zoo
"unsloth>=2026.6.7" unsloth-zoo
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2640,7 +2657,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.6.3" unsloth-zoo
--upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2694,7 +2711,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.6.3" --torch-backend=auto
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2905,7 +2922,10 @@ if [ -t 1 ]; then
case "${_reply:-y}" in
[Yy]*|"")
step "launch" "starting Unsloth Studio..."
"$VENV_DIR/bin/unsloth" studio -p 8888
# Detach stdin from the `curl | sh` pipe: as a foreground server the
# studio would otherwise drain the rest of this piped script, leaving
# the shell to die parsing the now-truncated tail (`unexpected fi`).
"$VENV_DIR/bin/unsloth" studio -p 8888 </dev/null
_LAUNCH_EXIT=$?
if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
echo ""

7
pyproject.toml
View file

@ -56,6 +56,7 @@ studio = [
"frontend/.git*",
"backend/requirements/**/*",
"backend/plugins/**/*",
"backend/assets/**/*.jinja",
"backend/core/data_recipe/oxc-validator/*.json",
"backend/core/data_recipe/oxc-validator/*.mjs",
]
@ -71,7 +72,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.6.3",
"unsloth_zoo>=2026.6.5",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -92,7 +93,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.6.3",
"unsloth_zoo>=2026.6.5",
"torchvision",
"unsloth[triton]",
]
@ -582,7 +583,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.6.3",
"unsloth_zoo>=2026.6.5",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

397
studio/backend/assets/chat_templates/gemma-4-edge.jinja Normal file
View file

@ -0,0 +1,397 @@
{#-
Gemma 4 chat template (E2B / E4B edge variant), vendored for Unsloth Studio.
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
Studio-local changes vs PR #118:
1. preserve_thinking defaults to false (see SETUP block below).
2. The empty "<|channel>thought\n<channel|>" block on enable_thinking=false is
NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it,
google/gemma-4-E4B-it) that omits it; only the 12b/26B-A4B/31B family emits it.
This file matches the E2B/E4B behavior; gemma-4.jinja keeps the larger-model one.
Applied to unsloth/gemma-4-E2B-it-GGUF and unsloth/gemma-4-E4B-it-GGUF so the
embedded GGUF template does not need re-downloading.
-#}
{%- macro format_parameters(properties, required, filter_keys=false) -%}
{%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in properties | dictsort -%}
{%- set add_comma = false -%}
{%- if not filter_keys or key not in standard_keys -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{{ key }}:{
{%- if value['description'] -%}
description:<|"|>{{ value['description'] }}<|"|>
{%- set add_comma = true -%}
{%- endif -%}
{%- if value['type'] | upper == 'STRING' -%}
{%- if value['enum'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
enum:{{ format_argument(value['enum']) }}
{%- endif -%}
{%- elif value['type'] | upper == 'ARRAY' -%}
{%- if value['items'] is mapping and value['items'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
items:{
{%- set ns_items = namespace(found_first=false) -%}
{%- for item_key, item_value in value['items'] | dictsort -%}
{%- if item_value is not none -%}
{%- if ns_items.found_first %},{% endif -%}
{%- set ns_items.found_first = true -%}
{%- if item_key == 'properties' -%}
properties:{
{%- if item_value is mapping -%}
{{- format_parameters(item_value, value['items']['required'] | default([])) -}}
{%- endif -%}
}
{%- elif item_key == 'required' -%}
required:[
{%- for req_item in item_value -%}
<|"|>{{- req_item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- elif item_key == 'type' -%}
{%- if item_value is string -%}
type:{{ format_argument(item_value | upper) }}
{%- else -%}
type:{{ format_argument(item_value | map('upper') | list) }}
{%- endif -%}
{%- else -%}
{{ item_key }}:{{ format_argument(item_value) }}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
}
{%- endif -%}
{%- endif -%}
{%- if value['nullable'] %}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
nullable:true
{%- endif -%}
{%- if value['type'] | upper == 'OBJECT' -%}
{%- if value['properties'] is defined and value['properties'] is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value['properties'], value['required'] | default([])) -}}
}
{%- elif value is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
}
{%- endif -%}
{%- if value['required'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
required:[
{%- for item in value['required'] | default([]) -%}
<|"|>{{- item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- endif -%}
{%- endif -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
type:<|"|>{{ value['type'] | upper }}<|"|>}
{%- endif -%}
{%- endfor -%}
{%- endmacro -%}
{%- macro format_function_declaration(tool_data) -%}
declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
{%- set params = tool_data['function']['parameters'] -%}
{%- if params -%}
,parameters:{
{%- if params['properties'] -%}
properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
{%- endif -%}
{%- if params['required'] -%}
required:[
{%- for item in params['required'] -%}
<|"|>{{- item -}}<|"|>
{{- ',' if not loop.last -}}
{%- endfor -%}
],
{%- endif -%}
{%- if params['type'] -%}
type:<|"|>{{- params['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
{%- if 'response' in tool_data['function'] -%}
{%- set response_declaration = tool_data['function']['response'] -%}
,response:{
{%- if response_declaration['description'] -%}
description:<|"|>{{- response_declaration['description'] -}}<|"|>,
{%- endif -%}
{%- if response_declaration['type'] | upper == 'OBJECT' -%}
type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
}
{%- endmacro -%}
{%- macro format_argument(argument, escape_keys=True) -%}
{%- if argument is none -%}
{{- 'null' -}}
{%- elif argument is string -%}
{{- '<|"|>' + argument + '<|"|>' -}}
{%- elif argument is boolean -%}
{{- 'true' if argument else 'false' -}}
{%- elif argument is mapping -%}
{{- '{' -}}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in argument | dictsort -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{%- if escape_keys -%}
{{- '<|"|>' + key + '<|"|>' -}}
{%- else -%}
{{- key -}}
{%- endif -%}
:{{- format_argument(value, escape_keys=escape_keys) -}}
{%- endfor -%}
{{- '}' -}}
{%- elif argument is sequence -%}
{{- '[' -}}
{%- for item in argument -%}
{{- format_argument(item, escape_keys=escape_keys) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- ']' -}}
{%- else -%}
{{- argument -}}
{%- endif -%}
{%- endmacro -%}
{%- macro strip_thinking(text) -%}
{%- set ns = namespace(result='') -%}
{%- for part in text.split('<channel|>') -%}
{%- if '<|channel>' in part -%}
{%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
{%- else -%}
{%- set ns.result = ns.result + part -%}
{%- endif -%}
{%- endfor -%}
{{- ns.result | trim -}}
{%- endmacro -%}
{%- macro format_tool_response_block(tool_name, response) -%}
{{- '<|tool_response>' -}}
{%- if response is mapping -%}
{{- 'response:' + tool_name + '{' -}}
{%- for key, value in response | dictsort -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- '}' -}}
{%- else -%}
{{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
{%- endif -%}
{{- '<tool_response|>' -}}
{%- endmacro -%}
{#- ===== SETUP ===== -#}
{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
{%- set loop_messages = messages -%}
{%- set enable_thinking = enable_thinking | default(false) -%}
{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#}
{%- set preserve_thinking = preserve_thinking | default(false) -%}
{{- bos_token -}}
{#- Handle System/Tool Definitions Block -#}
{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
{{- '<|turn>system\n' -}}
{#- Inject Thinking token at the very top of the FIRST system turn -#}
{%- if enable_thinking -%}
{{- '<|think|>\n' -}}
{%- set ns.prev_message_type = 'think' -%}
{%- endif -%}
{%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
{%- if messages[0]['content'] is string -%}
{{- messages[0]['content'] | trim -}}
{%- elif messages[0]['content'] is sequence -%}
{%- for item in messages[0]['content'] -%}
{{- item['text'] | trim + ' '-}}
{%- endfor -%}
{%- endif -%}
{%- set loop_messages = messages[1:] -%}
{%- endif -%}
{%- if tools -%}
{%- for tool in tools %}
{{- '<|tool>' -}}
{{- format_function_declaration(tool) | trim -}}
{{- '<tool|>' -}}
{%- endfor %}
{%- set ns.prev_message_type = 'tool' -%}
{%- endif -%}
{{- '<turn|>\n' -}}
{%- endif %}
{#- Pre-scan: find last user message index for reasoning guard -#}
{%- set ns_turn = namespace(last_user_idx=-1) -%}
{%- for i in range(loop_messages | length) -%}
{%- if loop_messages[i]['role'] == 'user' -%}
{%- set ns_turn.last_user_idx = i -%}
{%- endif -%}
{%- endfor -%}
{#- Loop through messages -#}
{%- for message in loop_messages -%}
{%- if message['role'] != 'tool' -%}
{%- set ns.prev_message_type = None -%}
{%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
{#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#}
{%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
{%- if not continue_same_model_turn -%}
{{- '<|turn>' + role + '\n' }}
{%- endif -%}
{#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#}
{%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
{%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%}
{%- if thinking_text and thinking_gate and message.get('tool_calls') -%}
{{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
{%- endif -%}
{%- if message.get('tool_calls') -%}
{%- for tool_call in message.get('tool_calls') -%}
{%- set function = tool_call['function'] -%}
{{- '<|tool_call>call:' + function['name'] + '{' -}}
{%- if function['arguments'] is mapping -%}
{%- set ns_args = namespace(found_first=false) -%}
{%- for key, value in function['arguments'] | dictsort -%}
{%- if ns_args.found_first %},{% endif -%}
{%- set ns_args.found_first = true -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- endfor -%}
{%- elif function['arguments'] is none -%}
{%- else -%}
{{- raise_exception(
"chat_template: tool_calls[].function.arguments must be a "
"JSON object (mapping), not a string. Deserialize arguments "
"before passing to the template."
) -}}
{%- endif -%}
{{- '}<tool_call|>' -}}
{%- endfor -%}
{%- set ns.prev_message_type = 'tool_call' -%}
{%- endif -%}
{%- set ns_tr_out = namespace(flag=false) -%}
{%- if message.get('tool_responses') -%}
{#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
{%- for tool_response in message.get('tool_responses') -%}
{{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endfor -%}
{%- elif message.get('tool_calls') -%}
{#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
{%- set ns_tool_scan = namespace(stopped=false) -%}
{%- for k in range(loop.index0 + 1, loop_messages | length) -%}
{%- if ns_tool_scan.stopped -%}
{%- elif loop_messages[k]['role'] != 'tool' -%}
{%- set ns_tool_scan.stopped = true -%}
{%- else -%}
{%- set follow = loop_messages[k] -%}
{#- Resolve tool_call_id to function name -#}
{%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
{%- for tc in message.get('tool_calls') -%}
{%- if tc.get('id') == follow.get('tool_call_id') -%}
{%- set ns_tname.name = tc['function']['name'] -%}
{%- endif -%}
{%- endfor -%}
{#- Handle content as string or content-parts array -#}
{%- set tool_body = follow.get('content') -%}
{%- if tool_body is string -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- elif tool_body is sequence and tool_body is not string -%}
{%- set ns_txt = namespace(s='') -%}
{%- for part in tool_body -%}
{%- if part.get('type') == 'text' -%}
{%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
{%- endif -%}
{%- endfor -%}
{{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
{%- for part in tool_body -%}
{%- if part.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif part.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif part.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- else -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- endif -%}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- set captured_content -%}
{%- if message.get('content') is string -%}
{%- if role == 'model' -%}
{{- strip_thinking(message['content']) -}}
{%- else -%}
{{- message['content'] | trim -}}
{%- endif -%}
{%- elif message.get('content') is sequence -%}
{%- for item in message['content'] -%}
{%- if item.get('type') == 'text' -%}
{%- if role == 'model' -%}
{{- strip_thinking(item['text']) -}}
{%- else -%}
{{- item['text'] | trim -}}
{%- endif -%}
{%- elif item.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif item.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif item.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- endset -%}
{{- captured_content -}}
{%- set has_content = captured_content | trim | length > 0 -%}
{#- Forward-scan: find next non-tool message role for continuation detection -#}
{%- set next_nt = namespace(role=None, found=false) -%}
{%- for j in range(loop.index0 + 1, loop_messages | length) -%}
{%- if not next_nt.found -%}
{%- if loop_messages[j]['role'] != 'tool' -%}
{%- set next_nt.role = loop_messages[j]['role'] -%}
{%- set next_nt.found = true -%}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{%- set continues_into_next = (
role == 'model'
and next_nt.role == 'assistant'
and not message.get('tool_calls')
and not ns_tr_out.flag
) -%}
{%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
{{- '<|tool_response>' -}}
{%- elif continues_into_next -%}
{{- '\n' -}}
{%- elif not (ns_tr_out.flag and not has_content) -%}
{{- '<turn|>\n' -}}
{%- endif -%}
{#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
{%- set ns.prev_non_tool_role = message['role'] -%}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
{{- '<|turn>model\n' -}}
{%- endif -%}
{#- E2B/E4B do NOT emit an empty thought block when enable_thinking is false
(unlike the 12b/26B-A4B/31B family); see header. -#}
{%- endif -%}

397
studio/backend/assets/chat_templates/gemma-4.jinja Normal file
View file

@ -0,0 +1,397 @@
{#-
Gemma 4 chat template, vendored for Unsloth Studio.
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
Studio-local change: preserve_thinking defaults to false (see SETUP block below).
Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not
need re-downloading. Keep in sync with upstream if PR #118 changes.
-#}
{%- macro format_parameters(properties, required, filter_keys=false) -%}
{%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in properties | dictsort -%}
{%- set add_comma = false -%}
{%- if not filter_keys or key not in standard_keys -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{{ key }}:{
{%- if value['description'] -%}
description:<|"|>{{ value['description'] }}<|"|>
{%- set add_comma = true -%}
{%- endif -%}
{%- if value['type'] | upper == 'STRING' -%}
{%- if value['enum'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
enum:{{ format_argument(value['enum']) }}
{%- endif -%}
{%- elif value['type'] | upper == 'ARRAY' -%}
{%- if value['items'] is mapping and value['items'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
items:{
{%- set ns_items = namespace(found_first=false) -%}
{%- for item_key, item_value in value['items'] | dictsort -%}
{%- if item_value is not none -%}
{%- if ns_items.found_first %},{% endif -%}
{%- set ns_items.found_first = true -%}
{%- if item_key == 'properties' -%}
properties:{
{%- if item_value is mapping -%}
{{- format_parameters(item_value, value['items']['required'] | default([])) -}}
{%- endif -%}
}
{%- elif item_key == 'required' -%}
required:[
{%- for req_item in item_value -%}
<|"|>{{- req_item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- elif item_key == 'type' -%}
{%- if item_value is string -%}
type:{{ format_argument(item_value | upper) }}
{%- else -%}
type:{{ format_argument(item_value | map('upper') | list) }}
{%- endif -%}
{%- else -%}
{{ item_key }}:{{ format_argument(item_value) }}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
}
{%- endif -%}
{%- endif -%}
{%- if value['nullable'] %}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
nullable:true
{%- endif -%}
{%- if value['type'] | upper == 'OBJECT' -%}
{%- if value['properties'] is defined and value['properties'] is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value['properties'], value['required'] | default([])) -}}
}
{%- elif value is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
}
{%- endif -%}
{%- if value['required'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
required:[
{%- for item in value['required'] | default([]) -%}
<|"|>{{- item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- endif -%}
{%- endif -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
type:<|"|>{{ value['type'] | upper }}<|"|>}
{%- endif -%}
{%- endfor -%}
{%- endmacro -%}
{%- macro format_function_declaration(tool_data) -%}
declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
{%- set params = tool_data['function']['parameters'] -%}
{%- if params -%}
,parameters:{
{%- if params['properties'] -%}
properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
{%- endif -%}
{%- if params['required'] -%}
required:[
{%- for item in params['required'] -%}
<|"|>{{- item -}}<|"|>
{{- ',' if not loop.last -}}
{%- endfor -%}
],
{%- endif -%}
{%- if params['type'] -%}
type:<|"|>{{- params['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
{%- if 'response' in tool_data['function'] -%}
{%- set response_declaration = tool_data['function']['response'] -%}
,response:{
{%- if response_declaration['description'] -%}
description:<|"|>{{- response_declaration['description'] -}}<|"|>,
{%- endif -%}
{%- if response_declaration['type'] | upper == 'OBJECT' -%}
type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
}
{%- endmacro -%}
{%- macro format_argument(argument, escape_keys=True) -%}
{%- if argument is none -%}
{{- 'null' -}}
{%- elif argument is string -%}
{{- '<|"|>' + argument + '<|"|>' -}}
{%- elif argument is boolean -%}
{{- 'true' if argument else 'false' -}}
{%- elif argument is mapping -%}
{{- '{' -}}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in argument | dictsort -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{%- if escape_keys -%}
{{- '<|"|>' + key + '<|"|>' -}}
{%- else -%}
{{- key -}}
{%- endif -%}
:{{- format_argument(value, escape_keys=escape_keys) -}}
{%- endfor -%}
{{- '}' -}}
{%- elif argument is sequence -%}
{{- '[' -}}
{%- for item in argument -%}
{{- format_argument(item, escape_keys=escape_keys) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- ']' -}}
{%- else -%}
{{- argument -}}
{%- endif -%}
{%- endmacro -%}
{%- macro strip_thinking(text) -%}
{%- set ns = namespace(result='') -%}
{%- for part in text.split('<channel|>') -%}
{%- if '<|channel>' in part -%}
{%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
{%- else -%}
{%- set ns.result = ns.result + part -%}
{%- endif -%}
{%- endfor -%}
{{- ns.result | trim -}}
{%- endmacro -%}
{%- macro format_tool_response_block(tool_name, response) -%}
{{- '<|tool_response>' -}}
{%- if response is mapping -%}
{{- 'response:' + tool_name + '{' -}}
{%- for key, value in response | dictsort -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- '}' -}}
{%- else -%}
{{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
{%- endif -%}
{{- '<tool_response|>' -}}
{%- endmacro -%}
{#- ===== SETUP ===== -#}
{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
{%- set loop_messages = messages -%}
{%- set enable_thinking = enable_thinking | default(false) -%}
{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#}
{%- set preserve_thinking = preserve_thinking | default(false) -%}
{{- bos_token -}}
{#- Handle System/Tool Definitions Block -#}
{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
{{- '<|turn>system\n' -}}
{#- Inject Thinking token at the very top of the FIRST system turn -#}
{%- if enable_thinking -%}
{{- '<|think|>\n' -}}
{%- set ns.prev_message_type = 'think' -%}
{%- endif -%}
{%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
{%- if messages[0]['content'] is string -%}
{{- messages[0]['content'] | trim -}}
{%- elif messages[0]['content'] is sequence -%}
{%- for item in messages[0]['content'] -%}
{{- item['text'] | trim + ' '-}}
{%- endfor -%}
{%- endif -%}
{%- set loop_messages = messages[1:] -%}
{%- endif -%}
{%- if tools -%}
{%- for tool in tools %}
{{- '<|tool>' -}}
{{- format_function_declaration(tool) | trim -}}
{{- '<tool|>' -}}
{%- endfor %}
{%- set ns.prev_message_type = 'tool' -%}
{%- endif -%}
{{- '<turn|>\n' -}}
{%- endif %}
{#- Pre-scan: find last user message index for reasoning guard -#}
{%- set ns_turn = namespace(last_user_idx=-1) -%}
{%- for i in range(loop_messages | length) -%}
{%- if loop_messages[i]['role'] == 'user' -%}
{%- set ns_turn.last_user_idx = i -%}
{%- endif -%}
{%- endfor -%}
{#- Loop through messages -#}
{%- for message in loop_messages -%}
{%- if message['role'] != 'tool' -%}
{%- set ns.prev_message_type = None -%}
{%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
{#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#}
{%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
{%- if not continue_same_model_turn -%}
{{- '<|turn>' + role + '\n' }}
{%- endif -%}
{#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#}
{%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
{%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%}
{%- if thinking_text and thinking_gate and message.get('tool_calls') -%}
{{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
{%- endif -%}
{%- if message.get('tool_calls') -%}
{%- for tool_call in message.get('tool_calls') -%}
{%- set function = tool_call['function'] -%}
{{- '<|tool_call>call:' + function['name'] + '{' -}}
{%- if function['arguments'] is mapping -%}
{%- set ns_args = namespace(found_first=false) -%}
{%- for key, value in function['arguments'] | dictsort -%}
{%- if ns_args.found_first %},{% endif -%}
{%- set ns_args.found_first = true -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- endfor -%}
{%- elif function['arguments'] is none -%}
{%- else -%}
{{- raise_exception(
"chat_template: tool_calls[].function.arguments must be a "
"JSON object (mapping), not a string. Deserialize arguments "
"before passing to the template."
) -}}
{%- endif -%}
{{- '}<tool_call|>' -}}
{%- endfor -%}
{%- set ns.prev_message_type = 'tool_call' -%}
{%- endif -%}
{%- set ns_tr_out = namespace(flag=false) -%}
{%- if message.get('tool_responses') -%}
{#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
{%- for tool_response in message.get('tool_responses') -%}
{{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endfor -%}
{%- elif message.get('tool_calls') -%}
{#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
{%- set ns_tool_scan = namespace(stopped=false) -%}
{%- for k in range(loop.index0 + 1, loop_messages | length) -%}
{%- if ns_tool_scan.stopped -%}
{%- elif loop_messages[k]['role'] != 'tool' -%}
{%- set ns_tool_scan.stopped = true -%}
{%- else -%}
{%- set follow = loop_messages[k] -%}
{#- Resolve tool_call_id to function name -#}
{%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
{%- for tc in message.get('tool_calls') -%}
{%- if tc.get('id') == follow.get('tool_call_id') -%}
{%- set ns_tname.name = tc['function']['name'] -%}
{%- endif -%}
{%- endfor -%}
{#- Handle content as string or content-parts array -#}
{%- set tool_body = follow.get('content') -%}
{%- if tool_body is string -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- elif tool_body is sequence and tool_body is not string -%}
{%- set ns_txt = namespace(s='') -%}
{%- for part in tool_body -%}
{%- if part.get('type') == 'text' -%}
{%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
{%- endif -%}
{%- endfor -%}
{{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
{%- for part in tool_body -%}
{%- if part.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif part.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif part.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- else -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- endif -%}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- set captured_content -%}
{%- if message.get('content') is string -%}
{%- if role == 'model' -%}
{{- strip_thinking(message['content']) -}}
{%- else -%}
{{- message['content'] | trim -}}
{%- endif -%}
{%- elif message.get('content') is sequence -%}
{%- for item in message['content'] -%}
{%- if item.get('type') == 'text' -%}
{%- if role == 'model' -%}
{{- strip_thinking(item['text']) -}}
{%- else -%}
{{- item['text'] | trim -}}
{%- endif -%}
{%- elif item.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif item.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif item.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- endset -%}
{{- captured_content -}}
{%- set has_content = captured_content | trim | length > 0 -%}
{#- Forward-scan: find next non-tool message role for continuation detection -#}
{%- set next_nt = namespace(role=None, found=false) -%}
{%- for j in range(loop.index0 + 1, loop_messages | length) -%}
{%- if not next_nt.found -%}
{%- if loop_messages[j]['role'] != 'tool' -%}
{%- set next_nt.role = loop_messages[j]['role'] -%}
{%- set next_nt.found = true -%}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{%- set continues_into_next = (
role == 'model'
and next_nt.role == 'assistant'
and not message.get('tool_calls')
and not ns_tr_out.flag
) -%}
{%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
{{- '<|tool_response>' -}}
{%- elif continues_into_next -%}
{{- '\n' -}}
{%- elif not (ns_tr_out.flag and not has_content) -%}
{{- '<turn|>\n' -}}
{%- endif -%}
{#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
{%- set ns.prev_non_tool_role = message['role'] -%}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
{{- '<|turn>model\n' -}}
{%- endif -%}
{%- if not enable_thinking -%}
{#- Suppress thinking - but not when awaiting tool responses -#}
{%- if ns.prev_message_type != 'tool_call' -%}
{{- '<|channel>thought\n<channel|>' -}}
{%- endif -%}
{%- endif -%}
{%- endif -%}

159
studio/backend/cloudflare_tunnel.py
View file

@ -24,12 +24,20 @@ from pathlib import Path
from typing import Optional, Tuple
# cloudflared logs the quick-tunnel URL; match only the URL so we do not depend
# on the surrounding wording, which Cloudflare may change.
_URL_RE = re.compile(r"https://[A-Za-z0-9-]+\.trycloudflare\.com")
# on the surrounding wording, which Cloudflare may change. The negative lookahead
# drops cloudflared's own API host, which appears in failure lines such as
# failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel"
# and must never be mistaken for a usable tunnel URL.
_URL_RE = re.compile(r"https://(?!api\.)[A-Za-z0-9-]+\.trycloudflare\.com")
# cloudflared logs this once per edge connection it establishes. Until at least
# one appears the quick-tunnel URL returns Cloudflare error 1033 (HTTP 530), so
# we wait for it before advertising the URL.
_REGISTERED_MARKER = "Registered tunnel connection"
_RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download"
_URL_TIMEOUT = 15.0 # seconds to wait for the public URL before giving up
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
@ -180,13 +188,24 @@ class CloudflareTunnel:
upstream stays local-only.
"""
def __init__(self, port: int, binary: str):
def __init__(
self,
port: int,
binary: str,
protocol: Optional[str] = None,
):
self.port = port
self.binary = binary
# None lets cloudflared pick its default (quic, with its own http2
# fallback); set to "http2" to force it when quic is blocked.
self.protocol = protocol
self._proc: Optional[subprocess.Popen] = None
self._lock = threading.Lock()
self._stopped = False
self._url_event = threading.Event()
self._ready_event = threading.Event()
self.url: Optional[str] = None
self.ready = False
self.error: Optional[str] = None
def start(self) -> None:
@ -197,25 +216,33 @@ class CloudflareTunnel:
f"http://localhost:{self.port}",
"--no-autoupdate",
]
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
stdin = subprocess.DEVNULL,
text = True,
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
)
if self.protocol:
cmd += ["--protocol", self.protocol]
with self._lock:
# A stop() that landed before us (e.g. a shutdown in the caller's
# register->start window) marks the tunnel stopped; spawning now would
# orphan a process nobody owns, so refuse.
if self._stopped:
return
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
stdin = subprocess.DEVNULL,
text = True,
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
)
self._proc = proc
threading.Thread(
target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True
).start()
def _reader(self, proc: subprocess.Popen) -> None:
# Drain cloudflared's output, capture the first trycloudflare URL, and
# keep draining so it never blocks on a full pipe.
# Drain cloudflared's output: capture the first trycloudflare URL and the
# first edge-connection registration, and keep draining so it never
# blocks on a full pipe.
try:
if proc.stdout is not None:
for line in proc.stdout:
@ -224,20 +251,35 @@ class CloudflareTunnel:
if match:
self.url = match.group(0)
self._url_event.set()
if not self.ready and _REGISTERED_MARKER in line:
self.ready = True
self._ready_event.set()
except Exception:
pass
finally:
# stdout closed -> cloudflared has exited. Record why, and unblock any
# waiters at once instead of letting them wait out the full timeout.
if self.url is None:
self.error = "cloudflared exited before emitting a tunnel URL"
self._url_event.set()
elif not self.ready:
self.error = "cloudflared exited before the tunnel connection registered"
self._url_event.set()
self._ready_event.set()
def wait_for_url(self, timeout: float = _URL_TIMEOUT) -> Optional[str]:
self._url_event.wait(timeout)
return self.url
def wait_for_ready(self, timeout: float = _READY_TIMEOUT) -> Optional[str]:
"""Block until the tunnel is actually serving -- the URL has been minted
*and* at least one edge connection has registered -- or until timeout.
Returns the URL only when ready, so callers never advertise a URL that
would return Cloudflare error 1033 (HTTP 530)."""
self._ready_event.wait(timeout)
return self.url if self.ready else None
def stop(self) -> None:
"""Terminate the tunnel. Idempotent and safe to call from a signal handler."""
with self._lock:
# Mark stopped so a start() racing behind us refuses to spawn.
self._stopped = True
proc, self._proc = self._proc, None
if proc is None:
return
@ -260,43 +302,74 @@ class CloudflareTunnel:
# enough; the lock guards the start/stop/shutdown races.
_active_tunnel: Optional[CloudflareTunnel] = None
_active_lock = threading.Lock()
# Latched by stop_studio_tunnel so a shutdown landing *between* a start's retry
# attempts aborts the loop instead of starting a tunnel nobody will ever stop.
_shutdown_requested = False
def start_studio_tunnel(port: int, timeout: float = _URL_TIMEOUT) -> Optional[str]:
"""Start a quick tunnel and return its public URL, or None (best-effort).
def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[str]:
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
On any failure (no binary, no URL within timeout, early crash) the tunnel is
stopped and None is returned, so the caller prints a hint and continues.
Waits for cloudflared to both mint the URL and register an edge connection
before returning, so the caller never advertises a URL that yields Cloudflare
error 1033 (HTTP 530). If a URL is minted but no connection registers within
the window (e.g. quic is blocked on this network), retries once forcing the
http2 protocol. On any failure the tunnel is stopped and None is returned.
"""
global _active_tunnel
global _active_tunnel, _shutdown_requested
binary = ensure_cloudflared()
if not binary:
return None
tunnel = CloudflareTunnel(port, binary)
# Register before start/wait so a shutdown during the URL wait can stop it.
with _active_lock:
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
try:
tunnel.start()
url = tunnel.wait_for_url(timeout)
except Exception:
url = None
if url:
return url
# No URL (or crash): drop it unless a concurrent shutdown already replaced it.
with _active_lock:
if _active_tunnel is tunnel:
_active_tunnel = None
tunnel.stop()
_shutdown_requested = False # fresh session
# Default protocol first (quic, with cloudflared's own http2 fallback); if a
# URL appears but no connection registers, quic is likely blocked -> retry
# once forcing http2.
for protocol in (None, "http2"):
# Create + register under the lock, and bail if a stop already landed
# (e.g. between this and the previous attempt) so we never start a tunnel
# after shutdown has run.
with _active_lock:
if _shutdown_requested:
_active_tunnel = None
return None
tunnel = CloudflareTunnel(port, binary, protocol = protocol)
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
except Exception:
url = None
if url:
return url
saw_url = tunnel.url is not None
# Not ready: drop it, but only if we are still the active tunnel.
with _active_lock:
was_active = _active_tunnel is tunnel
if was_active:
_active_tunnel = None
tunnel.stop()
# A concurrent shutdown or start took over while we waited; retrying would
# spawn a tunnel nobody owns (orphaned after shutdown), so bail instead.
if not was_active:
return None
# No URL at all is an API/network failure, not a protocol one; forcing
# http2 will not help, so do not burn another window on it.
if not saw_url:
return None
return None
def stop_studio_tunnel() -> None:
"""Terminate the active tunnel, if any. Idempotent."""
global _active_tunnel
global _active_tunnel, _shutdown_requested
with _active_lock:
# Latch so an in-flight start_studio_tunnel won't start a fresh tunnel
# (e.g. its http2 retry) after we have already torn down.
_shutdown_requested = True
tunnel, _active_tunnel = _active_tunnel, None
if tunnel is not None:
tunnel.stop()

37
studio/backend/core/export/export.py
View file

@ -18,7 +18,12 @@ from utils.hardware import clear_gpu_cache
from utils.models import is_vision_model, get_base_model_from_lora
from utils.models.model_config import detect_audio_type
from utils.paths import ensure_dir, outputs_root, resolve_export_dir, resolve_output_dir
from utils.paths import (
ensure_dir,
outputs_root,
resolve_export_write_dir,
resolve_output_dir,
)
from core.inference import get_inference_backend
# GPU-only imports — guarded for Apple Silicon where these aren't needed
@ -336,7 +341,7 @@ class ExportBackend:
save_method = "merged_16bit"
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
save_directory = str(resolve_export_write_dir(save_directory))
logger.info(f"Saving merged model locally to: {save_directory}")
ensure_dir(Path(save_directory))
@ -436,7 +441,7 @@ class ExportBackend:
output_path: Optional[str] = None
try:
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
save_directory = str(resolve_export_write_dir(save_directory))
logger.info(f"Saving base model locally to: {save_directory}")
ensure_dir(Path(save_directory))
@ -563,6 +568,7 @@ class ExportBackend:
return False, "No model loaded. Please select a checkpoint first.", None
output_path: Optional[str] = None
model_tmp_to_cleanup: Optional[str] = None
try:
# unsloth expects lowercase quant method
quant_method = quantization_method.lower()
@ -588,9 +594,8 @@ class ExportBackend:
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
# Absolute path so unsloth's relative-path internals resolve
# against the repo root cwd, not the export directory.
save_directory = str(resolve_export_write_dir(save_directory))
# Keep unsloth relative-path internals anchored to the repo cwd.
abs_save_dir = os.path.abspath(save_directory)
logger.info(f"Saving GGUF model locally to: {abs_save_dir}")
@ -604,9 +609,15 @@ class ExportBackend:
cwd = os.getcwd()
pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf")))
model_save_path = os.path.join(abs_save_dir, "model")
pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()}
# Avoid clobbering an existing user-owned model/ directory.
import uuid
_model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}")
model_tmp_to_cleanup = _model_tmp
self.current_model.save_pretrained_gguf(
model_save_path,
_model_tmp,
self.current_tokenizer,
quantization_method = quant_method,
)
@ -618,10 +629,12 @@ class ExportBackend:
shutil.move(src, dest)
logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/")
# Flatten any .gguf from subdirs (e.g. model_gguf/) into abs_save_dir.
# Flatten GGUF files from subdirs created during this export.
for sub in list(Path(abs_save_dir).iterdir()):
if not sub.is_dir():
continue
if sub.name in pre_existing_subs:
continue
for src in sub.glob("*.gguf"):
dest = os.path.join(abs_save_dir, src.name)
shutil.move(str(src), dest)
@ -634,7 +647,7 @@ class ExportBackend:
if self.current_checkpoint:
ckpt = Path(self.current_checkpoint)
gguf_dir = ckpt.parent / f"{ckpt.name}_gguf"
if gguf_dir.is_dir():
if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve():
for src in gguf_dir.glob("*.gguf"):
dest = os.path.join(abs_save_dir, src.name)
shutil.move(str(src), dest)
@ -683,6 +696,8 @@ class ExportBackend:
)
except Exception as e:
if model_tmp_to_cleanup:
shutil.rmtree(model_tmp_to_cleanup, ignore_errors = True)
logger.error(f"Error exporting GGUF model: {e}")
import traceback
@ -712,7 +727,7 @@ class ExportBackend:
output_path: Optional[str] = None
try:
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
save_directory = str(resolve_export_write_dir(save_directory))
logger.info(f"Saving LoRA adapter locally to: {save_directory}")
ensure_dir(Path(save_directory))

109
studio/backend/core/inference/chat_templates.py Normal file
View file

@ -0,0 +1,109 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Bundled chat-template selection for GGUF inference.
Some shipped GGUF quants embed an older chat template. Rather than re-cutting and
asking users to re-download every quant, Studio can override the embedded template
at llama-server launch time with a bundled, up-to-date Jinja template for known
model families. The override is wired through the existing ``chat_template_override``
-> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``.
Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118
``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking"
toggle appears while staying disabled by default.
"""
import re
from functools import lru_cache
from pathlib import Path
from typing import Optional
# assets live at <backend>/assets/chat_templates/. This module is at
# <backend>/core/inference/chat_templates.py, so walk up three parents to
# <backend> (mirrors utils/inference/inference_config.py).
_ASSETS_DIR = Path(__file__).parent.parent.parent / "assets" / "chat_templates"
# unsloth/gemma-4-<variant>-GGUF (case-insensitive). The "-GGUF" suffix is retained
# on ModelConfig.identifier for HF GGUF repos, so this matches E2B / E4B / 31B /
# 26B-A4B and any future unsloth/gemma-4-*-GGUF, while excluding gemma-3,
# non-Unsloth, and non-GGUF identifiers (e.g. the bf16 "unsloth/gemma-4-E2B-it").
_GEMMA4_GGUF_RE = re.compile(r"^unsloth/gemma-4-.+-gguf$", re.IGNORECASE)
# Google ships two distinct gemma-4 chat templates: E2B/E4B omit the empty
# "<|channel>thought<channel|>" block on enable_thinking=false, while the
# 12b/26B-A4B/31B family emits it. Route the two GGUF families to the matching
# bundled template so each keeps its model's intended behavior.
_GEMMA4_EDGE_GGUF_RE = re.compile(r"^unsloth/gemma-4-e[24]b-it-gguf$", re.IGNORECASE)
_GEMMA4_TEMPLATE_FILE = "gemma-4.jinja" # 12b / 26B-A4B / 31B
_GEMMA4_EDGE_TEMPLATE_FILE = "gemma-4-edge.jinja" # E2B / E4B
def _canonical_repo_id(model_identifier: str) -> str:
"""Mirror ``ModelConfig.from_identifier``: a bare HF shorthand with no owner
(e.g. ``gemma-4-E2B-it-GGUF``) defaults to the ``unsloth/`` org. The resolver
runs on the raw ``request.model_path`` (before that canonicalization), so apply
the same rule here, otherwise shorthand loads would skip the override.
"""
mid = model_identifier.strip()
if mid and "/" not in mid:
mid = f"unsloth/{mid}"
return mid
def is_unsloth_gemma4_gguf(model_identifier: Optional[str]) -> bool:
"""True for canonical ``unsloth/gemma-4-*-GGUF`` repo identifiers (and the
owner-less shorthand that resolves to the same Unsloth repo)."""
if not model_identifier:
return False
return bool(_GEMMA4_GGUF_RE.match(_canonical_repo_id(model_identifier)))
def is_unsloth_gemma4_edge_gguf(model_identifier: Optional[str]) -> bool:
"""True for the E2B / E4B GGUF repos, which use the edge-variant template."""
if not model_identifier:
return False
return bool(_GEMMA4_EDGE_GGUF_RE.match(_canonical_repo_id(model_identifier)))
def _gemma4_template_file(model_identifier: Optional[str]) -> Optional[str]:
"""Return the bundled template filename for a gemma-4 GGUF id, else None."""
if is_unsloth_gemma4_edge_gguf(model_identifier):
return _GEMMA4_EDGE_TEMPLATE_FILE
if is_unsloth_gemma4_gguf(model_identifier):
return _GEMMA4_TEMPLATE_FILE
return None
@lru_cache(maxsize=8)
def load_bundled_chat_template(name: str) -> str:
"""Read a bundled chat-template asset by filename (cached for the process)."""
return (_ASSETS_DIR / name).read_text(encoding="utf-8")
def resolve_effective_chat_template_override(
*,
model_identifier: Optional[str],
user_override: Optional[str],
) -> Optional[str]:
"""Resolve which chat-template text to launch llama-server with.
Precedence:
1. An explicit, non-empty user override always wins (advanced users).
2. For ``unsloth/gemma-4-*-GGUF``, return the bundled gemma-4 template
(adds ``preserve_thinking``, default off) so the embedded GGUF template
is overridden without re-downloading quants. E2B/E4B get the edge
variant; 12b/26B-A4B/31B get the standard one.
3. Otherwise ``None`` -> llama-server renders the GGUF's embedded template.
The result is fed to ``LlamaCppBackend.load_model(chat_template_override=...)``
and must be computed before the route-level reload-dedup check so the live
backend state and the incoming request compare consistently.
"""
if user_override and user_override.strip():
return user_override
template_file = _gemma4_template_file(model_identifier)
if template_file is not None:
return load_bundled_chat_template(template_file)
return None

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

File diff suppressed because it is too large Load diff

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

@ -109,6 +109,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
out.append(token)
parse_ctx_override(out)
parse_cache_override(out)
parse_split_mode_override(out)
return out
@ -157,8 +158,20 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
"--no-jinja",
}
)
# Multi-GPU split mode shadows the Tensor Parallelism toggle
# (--split-mode tensor). Pass-through stays allowed so users keep the
# row/none/layer modes the toggle doesn't expose, but it's stripped on
# inherit and reconciled into the round-tripped tensor_parallel state.
# --tensor-split is coupled to the split mode and is stripped with it: Studio
# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must
# not last-wins-override Studio's computed asymmetric split.
_SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"})
_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"})
_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
_SHADOWING_FLAGS: frozenset[str] = _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS
)
# Shadowing flags that take no value -- strip the flag only, not the next token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"})
@ -213,11 +226,12 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) ->
return override if override is not None else fallback_n_ctx
def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last-wins cache type if extras pass cache flags.
def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]:
"""Return the last-wins string value among ``flags`` in extras, or None.
Recognises -ctk (key) and -ctv (value); treats both as one setting,
since Studio's KV estimate has a single cache_type_kv knob.
Handles both ``--flag=value`` and ``--flag value`` forms and raises if a
matched flag has no (or an empty) value. Shared by the single-knob
last-wins parsers (cache type, split mode).
"""
if not args:
return None
@ -228,7 +242,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
while i < n:
tok = tokens[i]
flag = _flag_name(tok)
if flag is None or flag not in _CACHE_FLAGS:
if flag is None or flag not in flags:
i += 1
continue
@ -249,6 +263,17 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
return override
def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last-wins cache type if extras pass cache flags.
Mirrors parse_ctx_override but for cache type. Recognises both -ctk
(key) and -ctv (value). When both flags appear, returns the last-wins
value, treating key and value cache flags as the same setting because
Studio's KV estimate has a single cache_type_kv knob.
"""
return _last_flag_value(args, _CACHE_FLAGS)
def resolve_cache_type_kv(
args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str]
) -> Optional[str]:
@ -260,6 +285,52 @@ def resolve_cache_type_kv(
return override if override is not None else fallback_cache_type_kv
def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last-wins ``--split-mode`` / ``-sm`` value from extras.
Mirrors parse_cache_override for the multi-GPU split mode. Returns the
raw mode string (e.g. ``tensor`` / ``row`` / ``none`` / ``layer``), or
None when extras don't set it.
"""
return _last_flag_value(args, _SPLIT_MODE_FLAGS)
def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool:
"""Return the tensor-parallel state load_model should treat as requested.
A user-supplied ``--split-mode`` in extras last-wins-overrides the
toggle, so reconcile it back into the boolean: any explicit split mode
means tensor-parallel is on iff that mode is ``tensor``. Falls back to
the toggle value when extras don't set it.
"""
override = parse_split_mode_override(args)
if override is None:
return fallback_tensor_parallel
return override.strip().lower() == "tensor"
_MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"})
_MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"})
def extra_args_disable_mmproj(args: Optional[Iterable[str]]) -> bool:
"""True when pass-through args opt out of vision mmproj loading.
llama-server parses --mmproj-auto / --no-mmproj / --no-mmproj-auto as one
boolean with last-wins semantics; mirror that here.
"""
if not args:
return False
disabled = False
for raw in args:
flag = _flag_name(str(raw))
if flag in _MMPROJ_DISABLE_FLAGS:
disabled = True
elif flag in _MMPROJ_ENABLE_FLAGS:
disabled = False
return disabled
def strip_shadowing_flags(
args: Iterable[str],
*,
@ -267,12 +338,15 @@ def strip_shadowing_flags(
strip_cache: bool = True,
strip_spec: bool = True,
strip_template: bool = True,
strip_split_mode: bool = True,
) -> list[str]:
"""Strip flags that shadow first-class Studio settings.
Used when inheriting a previous load's ``llama_extra_args`` so an
inherited `-c 4096` can't override the current `max_seq_length` (same for
cache / spec / template). Each ``strip_*`` toggle controls one group.
inherited `-c 4096` can't override the current `max_seq_length`
(same for cache / spec / template / split-mode). Each ``strip_*``
toggle controls one group; the route only strips groups whose
first-class field the caller actually supplied.
"""
shadowing: set[str] = set()
if strip_context:
@ -283,6 +357,8 @@ def strip_shadowing_flags(
shadowing |= _SPEC_FLAGS
if strip_template:
shadowing |= _TEMPLATE_FLAGS
if strip_split_mode:
shadowing |= _SPLIT_SHADOWING_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []
@ -303,3 +379,20 @@ def strip_shadowing_flags(
else:
i += 1
return out
def strip_split_mode_only(args: Optional[Iterable[str]]) -> Optional[list[str]]:
"""Remove the split-mode group (``--split-mode`` / ``-sm`` and the coupled
``--tensor-split`` / ``-ts``) from ``args``, keeping every other shadow flag.
Preserves a None/empty input so the inherit-vs-explicit-empty distinction
survives. Used where tensor mode is being forced off (downgrade / fallback)."""
if not args:
return args
return strip_shadowing_flags(
args,
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = True,
)

79
studio/backend/core/inference/mcp_client.py
View file

@ -8,6 +8,7 @@ import json
import os
import shlex
import sys
import time
from typing import Any, Optional
from loggers import get_logger
@ -16,6 +17,14 @@ logger = get_logger(__name__)
MCP_TOOL_PREFIX = "mcp__"
# A failed probe isn't cached (a recovered server must come back), but it's
# recorded so a down server isn't re-probed -- and the chat send re-hung for
# the full timeout -- on every message. Cool off for this long after a failure;
# much longer for OAuth, whose probe can hang up to _OAUTH_PROBE_TIMEOUT,
# so that hang doesn't recur every minute.
FAILED_PROBE_COOLOFF_SECONDS = 60.0
OAUTH_FAILED_PROBE_COOLOFF_SECONDS = 300.0
_oauth_token_store = None
@ -108,10 +117,25 @@ def join_stdio_command(parts: list[str]) -> str:
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (bypassing the
sandbox), so allowed only when the host is the user's own machine. The Tauri
app sets UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1; localhost/self-hosted users can opt
in with the same var. Off for Colab and any network (0.0.0.0) bind."""
return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1"
sandbox), so allowed only when the host is the user's own machine. On startup
a loopback bind defaults UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see
utils.host_policy.apply_stdio_mcp_loopback_default, called from run.py); the
Tauri app does the same. Off for Colab and any network (0.0.0.0) bind unless
an operator sets the var out-of-band; set it to 0 to force-disable.
When stdio is on only because of that loopback auto-default, an explicit
`unsloth studio run --disable-tools` turns it back off (a local stdio command
is server-side code execution). An explicit operator opt-in via the env var
still wins -- including the documented `=1` network opt-in, where the process
tool policy is False merely by the external-host default, not by choice."""
if os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") != "1":
return False
from state.tool_policy import get_tool_policy
from utils.host_policy import loopback_default_active
if loopback_default_active() and get_tool_policy() is False:
return False
return True
# Probe timeouts for discovering a server's tool list. OAuth needs minutes for
@ -227,6 +251,53 @@ async def list_tools_async(
return await asyncio.wait_for(_fetch(), timeout = timeout)
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools()
# probes a server only on a cache miss, keeping MCP discovery off the chat
# send's critical path -- tool schemas are stable within a session. The
# /refresh route warms it; a URL/header/OAuth change or a delete evicts it.
# Successful probes are cached indefinitely.
_tool_cache: dict[str, list[dict]] = {}
# server_id -> monotonic time before which a failed server must not be
# re-probed (see record_probe_failure). Cleared on a successful probe or
# eviction.
_probe_cooloff_until: dict[str, float] = {}
# MCP server fields whose change invalidates a server's discovered tools: the
# endpoint/auth used to probe it (url, headers, oauth) or whether it's used at
# all (is_enabled). A rename does not. The update route's eviction and
# get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift.
TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"})
def get_cached_tools(server_id: str) -> Optional[list[dict]]:
return _tool_cache.get(server_id)
def cache_tools(server_id: str, tools: list[dict]) -> None:
_tool_cache[server_id] = tools
_probe_cooloff_until.pop(server_id, None)
def record_probe_failure(server_id: str, use_oauth: bool = False) -> None:
cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS
_probe_cooloff_until[server_id] = time.monotonic() + cooloff
def in_failure_cooloff(server_id: str) -> bool:
return _probe_cooloff_until.get(server_id, 0.0) > time.monotonic()
def invalidate_tool_cache(server_id: Optional[str] = None) -> None:
"""Evict one server's cached tools, or every entry when server_id is None."""
if server_id is None:
_tool_cache.clear()
_probe_cooloff_until.clear()
else:
_tool_cache.pop(server_id, None)
_probe_cooloff_until.pop(server_id, None)
def _flatten_result(result: Any) -> str:
parts = []
for block in getattr(result, "content", None) or []:

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

@ -861,6 +861,7 @@ class InferenceOrchestrator:
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
**_unused,
@ -922,6 +923,7 @@ class InferenceOrchestrator:
tool_call_timeout = tool_call_timeout,
session_id = session_id,
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
)
def generate_with_adapter_control(

17
studio/backend/core/inference/providers.py
View file

@ -249,6 +249,23 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
# Surfaced via the frontend's CUSTOM_PROVIDER_PRESETS, not the dropdown.
"hidden": True,
},
"custom": {
"display_name": "Custom",
# User-supplied via provider_base_url.
"base_url": "",
"default_models": [],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": (
"User-supplied OpenAI-compatible server. Routed to "
"/v1/chat/completions; /models is optional."
),
# Surfaced by the frontend's generic Custom option, not the dropdown.
"hidden": True,
},
"ollama": {
"display_name": "Ollama",
"base_url": "http://localhost:11434/v1",

53
studio/backend/core/inference/safetensors_agentic.py
View file

@ -35,6 +35,13 @@ from core.inference.tool_loop_controller import (
status_for_tool,
tool_event_provenance,
)
from state.tool_approvals import (
TOOL_REJECTED_MESSAGE,
abort_tool_decision,
begin_tool_decision,
new_approval_id,
wait_tool_decision,
)
logger = get_logger(__name__)
@ -146,6 +153,7 @@ def run_safetensors_tool_loop(
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
@ -174,7 +182,7 @@ def run_safetensors_tool_loop(
# Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search.
from core.inference.tools import build_rag_autoinject
_auto = build_rag_autoinject(conversation, rag_scope)
_auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope)
if _auto:
for _ev in _auto["events"]:
yield _ev
@ -509,8 +517,47 @@ def run_safetensors_tool_loop(
else:
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
yield {"type": "status", "text": decision.status_text}
yield decision.tool_start_event()
needs_confirm = bool(confirm_tool_calls)
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
start_event = decision.tool_start_event()
start_event["approval_id"] = approval_id
start_event["awaiting_confirmation"] = needs_confirm
try:
yield {"type": "status", "text": decision.status_text}
yield start_event
if (
decision_slot is not None
and wait_tool_decision(
decision_slot,
approval_id,
cancel_event = cancel_event,
)
== "deny"
):
decision_slot = None
yield {
"type": "tool_end",
"tool_name": decision.tool_name,
"tool_call_id": decision.tool_call_id,
"result": TOOL_REJECTED_MESSAGE,
"provenance": decision.provenance,
}
denied_message = {
"role": "tool",
"name": decision.tool_name,
"content": TOOL_REJECTED_MESSAGE,
}
if decision.tool_call_id:
denied_message["tool_call_id"] = decision.tool_call_id
conversation.append(denied_message)
continue
decision_slot = None
finally:
if decision_slot is not None:
abort_tool_decision(decision_slot, approval_id)
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
# RAG: cap paraphrased KB re-searches that slip past the dup guard.

70
studio/backend/core/inference/tensor_fallback.py Normal file
View file

@ -0,0 +1,70 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tensor-parallel -> layer-split auto-fallback for GGUF loads.
Kept in its own module (no FastAPI / httpx deps) so the orchestration can be
unit-tested with a fake loader, without a GPU or a running llama-server.
"""
from __future__ import annotations
import logging
from typing import Awaitable, Callable, Optional
from core.inference.llama_server_args import (
resolve_tensor_parallel,
strip_split_mode_only,
)
logger = logging.getLogger(__name__)
async def load_with_tensor_fallback(
attempt_load: Callable[[bool, Optional[list[str]]], Awaitable[bool]],
*,
requested_tensor: bool,
extra_args: Optional[list[str]],
label: str = "",
cancelled: Optional[Callable[[], bool]] = None,
) -> bool:
"""Run a GGUF load with the tensor-parallel -> layer-split auto-fallback.
``attempt_load(tensor_parallel, extra_args)`` performs one load and returns
True on success; it *raises* on a hard crash (llama-server aborts on some
archs / older builds), which is treated the same as a False return.
Tensor mode can be requested by the toggle or by a ``--split-mode tensor``
in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether
tensor mode is actually engaged, and it strips ``--split-mode`` from the
extras so the layer retry can't relaunch the same failing tensor load. A
non-tensor load keeps its original contract and propagates exceptions.
``cancelled()`` distinguishes a real tensor-start failure from a user
cancellation: ``attempt_load`` also returns False when the load was
cancelled, so without this the helper would restart a load the user just
cancelled.
"""
tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor)
try:
success = await attempt_load(requested_tensor, extra_args)
except Exception as exc:
if not tensor_requested:
raise
logger.warning("Tensor-parallel load raised for '%s': %s", label, exc)
success = False
if success or not tensor_requested:
return success
# The first attempt returned False because the user cancelled, not because
# tensor mode is unsupported -- do not relaunch the cancelled load.
if cancelled is not None and cancelled():
return success
logger.warning(
"Tensor-parallel load failed for '%s'; retrying with layer split "
"(this model may not support tensor parallelism)",
label,
)
return await attempt_load(False, strip_split_mode_only(extra_args))

75
studio/backend/core/inference/tools.py
View file

@ -24,11 +24,16 @@ import urllib.request
from core.inference.mcp_client import (
MCP_TOOL_PREFIX,
TOOL_CACHE_INVALIDATING_FIELDS,
cache_tools,
call_tool_sync,
get_cached_tools,
in_failure_cooloff,
is_stdio,
list_tools_async,
parse_server_headers,
probe_timeout,
record_probe_failure,
stdio_mcp_enabled,
)
from storage import mcp_servers_db
@ -634,28 +639,56 @@ async def get_enabled_mcp_tools() -> list[dict]:
if not servers:
return []
results = await asyncio.gather(
*(
list_tools_async(
url = s["url"],
headers = parse_server_headers(s),
timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))),
use_oauth = bool(s.get("use_oauth")),
)
for s in servers
),
return_exceptions = True,
)
# Skip servers still in their post-failure cool-off, otherwise a down
# server gets re-probed -- and blocks the send for the full timeout -- on
# every message.
uncached = [
s for s in servers if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"])
]
if uncached:
results = await asyncio.gather(
*(
list_tools_async(
url = s["url"],
headers = parse_server_headers(s),
timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))),
use_oauth = bool(s.get("use_oauth")),
)
for s in uncached
),
return_exceptions = True,
)
# An edit/delete can land while we await a probe (up to 305 s for
# OAuth); its cache eviction is a no-op against an entry we haven't
# written yet. Re-read and drop a result whose server changed or
# was removed mid-probe, else a stale tool list caches indefinitely.
current = {s["id"]: s for s in mcp_servers_db.list_servers()}
for server, payload in zip(uncached, results):
# Guard the failure branch too: a stale failure must not park a
# cool-off on the fresh config, or the server the user just fixed
# is skipped for the whole window.
fresh = current.get(server["id"])
if fresh is None or any(
fresh.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS
):
continue
if isinstance(payload, BaseException):
logger.warning(
"MCP server '%s' (%s) discovery failed: %s",
server.get("display_name") or server["id"],
server.get("url"),
payload,
)
# Failures aren't cached, but record one so a down server
# isn't re-probed every send during the cool-off.
record_probe_failure(server["id"], bool(fresh.get("use_oauth")))
continue
cache_tools(server["id"], payload)
specs: list[dict] = []
for server, payload in zip(servers, results):
if isinstance(payload, BaseException):
logger.warning(
"MCP server '%s' (%s) discovery failed: %s",
server.get("display_name") or server["id"],
server.get("url"),
payload,
)
for server in servers:
payload = get_cached_tools(server["id"])
if payload is None:
continue
specs.extend(_mcp_specs_for_server(server, payload))
return specs
@ -773,6 +806,7 @@ def _search_knowledge_base(arguments: dict, rag_scope: dict | None) -> str:
query = str(query),
scope_kb_id = scope.get("kb_id"),
scope_thread_id = scope.get("thread_id"),
scope_project_id = scope.get("project_id"),
top_k = top_k,
**_scope_retrieval_kwargs(scope),
)
@ -880,6 +914,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
query = query,
scope_kb_id = rag_scope.get("kb_id"),
scope_thread_id = rag_scope.get("thread_id"),
scope_project_id = rag_scope.get("project_id"),
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),

22
studio/backend/core/rag/ingestion.py
View file

@ -35,6 +35,22 @@ def _sha256_file(path: str) -> str:
return h.hexdigest()
def _remove_upload(stored_path: str | None, *, keep_path: str | None = None) -> None:
if not stored_path:
return
try:
target = os.path.realpath(stored_path)
if keep_path is not None and target == os.path.realpath(keep_path):
return
from utils.paths import rag_uploads_root
uploads = os.path.realpath(str(rag_uploads_root()))
if os.path.isfile(target) and os.path.commonpath([uploads, target]) == uploads:
os.remove(target)
except Exception: # noqa: BLE001 - upload cleanup must not block ingestion.
logger.warning("failed to remove RAG upload %s", stored_path, exc_info = True)
def _emit(job_id: str, event: dict) -> None:
with _jobs_lock:
q = _jobs.get(job_id)
@ -152,6 +168,7 @@ def start_ingestion(
filename: str,
stored_path: str,
*,
project_id: str | None = None,
model_name: str | None = None,
) -> tuple[str, str]:
"""Create the document + job rows and spawn the worker, returning
@ -167,11 +184,15 @@ def start_ingestion(
existing = store.document_by_hash(conn, scope, sha)
if existing is not None:
job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0)
_remove_upload(stored_path)
with _jobs_lock:
_jobs[job_id] = queue.Queue()
_emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True})
_emit(job_id, None)
return existing, job_id
for failed in store.failed_documents_by_hash(conn, scope, sha):
store.delete_document(conn, failed["id"])
_remove_upload(failed.get("stored_path"), keep_path = stored_path)
document_id = store.create_document(
conn,
@ -180,6 +201,7 @@ def start_ingestion(
sha256 = sha,
kb_id = kb_id,
thread_id = thread_id,
project_id = project_id,
status = "pending",
stored_path = stored_path,
)

6
studio/backend/core/rag/retrieval.py
View file

@ -22,7 +22,7 @@ class Hit:
def retrieve_lexical(
conn: sqlite3.Connection,
scope: str,
scope: str | list[str],
query: str,
k: int | None = None,
) -> list[Hit]:
@ -32,7 +32,7 @@ def retrieve_lexical(
def retrieve_dense(
conn: sqlite3.Connection,
scope: str,
scope: str | list[str],
query: str,
k: int | None = None,
*,
@ -69,7 +69,7 @@ def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]:
def retrieve_hybrid(
conn: sqlite3.Connection,
scope: str,
scope: str | list[str],
query: str,
*,
k: int | None = None,

67
studio/backend/core/rag/store.py
View file

@ -29,6 +29,15 @@ def thread_scope(thread_id: str) -> str:
return f"thread_{thread_id}"
def project_scope(project_id: str) -> str:
return f"project_{project_id}"
def _scopes(scope) -> list[str]:
"""Search helpers accept one scope or several (e.g. project + thread)."""
return [scope] if isinstance(scope, str) else list(scope)
def _f32(vector) -> bytes:
"""Pack a vector into float32 bytes for vec0."""
return struct.pack(f"{len(vector)}f", *(float(x) for x in vector))
@ -96,19 +105,21 @@ def create_document(
sha256: str,
kb_id: str | None = None,
thread_id: str | None = None,
project_id: str | None = None,
status: str = "pending",
stored_path: str | None = None,
document_id: str | None = None,
) -> str:
document_id = document_id or str(uuid.uuid4())
conn.execute(
"INSERT INTO documents(id, scope, kb_id, thread_id, filename, sha256, status, "
"stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?)",
"INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, "
"status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)",
(
document_id,
scope,
kb_id,
thread_id,
project_id,
filename,
sha256,
status,
@ -137,7 +148,8 @@ def set_document_status(
def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]:
rows = conn.execute(
"SELECT id, scope, kb_id, thread_id, filename, sha256, status, error, num_chunks, created_at "
"SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, "
"num_chunks, created_at "
"FROM documents WHERE scope=? ORDER BY created_at DESC",
(scope,),
).fetchall()
@ -151,11 +163,21 @@ def get_document(conn: sqlite3.Connection, document_id: str) -> dict | None:
def document_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> str | None:
row = conn.execute(
"SELECT id FROM documents WHERE scope=? AND sha256=?", (scope, sha256)
"SELECT id FROM documents WHERE scope=? AND sha256=? AND status!='failed' "
"ORDER BY created_at DESC LIMIT 1",
(scope, sha256),
).fetchone()
return row["id"] if row else None
def failed_documents_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> list[dict]:
rows = conn.execute(
"SELECT id, stored_path FROM documents WHERE scope=? AND sha256=? AND status='failed'",
(scope, sha256),
).fetchall()
return [dict(r) for r in rows]
def add_chunks(
conn: sqlite3.Connection,
scope: str,
@ -220,30 +242,41 @@ def delete_document(conn: sqlite3.Connection, document_id: str) -> None:
conn.commit()
def search_lexical(conn: sqlite3.Connection, scope: str, query: str, k: int):
"""BM25 lexical search. Returns [(chunk_id, score)], higher = better."""
def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int):
"""BM25 lexical search over one scope or several. Returns
[(chunk_id, score)], higher = better."""
mq = _match_query(query)
if not mq:
return []
scopes = _scopes(scope)
if not scopes:
return []
placeholders = ",".join("?" * len(scopes))
rows = conn.execute(
"SELECT chunk_id, bm25(chunks_fts) AS s FROM chunks_fts "
"WHERE chunks_fts MATCH ? AND scope=? ORDER BY s LIMIT ?",
(mq, scope, k),
f"SELECT chunk_id, bm25(chunks_fts) AS s FROM chunks_fts "
f"WHERE chunks_fts MATCH ? AND scope IN ({placeholders}) ORDER BY s LIMIT ?",
(mq, *scopes, k),
).fetchall()
# bm25() is negative (more negative = better); flip to higher-is-better.
return [(r["chunk_id"], -r["s"]) for r in rows]
def search_dense(conn: sqlite3.Connection, scope: str, vector, k: int):
"""Cosine KNN over vec0. Returns [(chunk_id, 1 - distance)]."""
def search_dense(conn: sqlite3.Connection, scope, vector, k: int):
"""Cosine KNN over vec0 for one scope or several. Returns
[(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by
equality, so multi-scope runs one query per scope and merges by score."""
if not rag_db.vec_table_exists(conn):
return []
rows = conn.execute(
"SELECT chunk_id, distance FROM chunks_vec "
"WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?",
(scope, _f32(vector), k),
).fetchall()
return [(r["chunk_id"], 1.0 - r["distance"]) for r in rows]
out: list[tuple[str, float]] = []
for s in _scopes(scope):
rows = conn.execute(
"SELECT chunk_id, distance FROM chunks_vec "
"WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?",
(s, _f32(vector), k),
).fetchall()
out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows)
out.sort(key = lambda t: t[1], reverse = True)
return out[:k]
def chunks_by_id(conn: sqlite3.Connection, ids) -> dict:

30
studio/backend/core/rag/tool.py
View file

@ -3,7 +3,8 @@
"""``search_knowledge_base`` LLM tool: scope resolution + hit formatting.
KB scope wins over thread scope. Hits render as ``<chunk>`` blocks for the model,
KB scope wins; otherwise project and thread scopes combine so project chats also
see their own attachments. Hits render as ``<chunk>`` blocks for the model,
plus a parallel citation source-map for clickable sources. Each call opens and
closes its own ``rag_db`` connection.
"""
@ -15,7 +16,7 @@ from xml.sax.saxutils import quoteattr
from storage import rag_db
from . import config, retrieval
from .store import kb_scope, thread_scope
from .store import kb_scope, project_scope, thread_scope
SEARCH_KNOWLEDGE_BASE_TOOL = {
"type": "function",
@ -42,12 +43,23 @@ SEARCH_KNOWLEDGE_BASE_TOOL = {
}
def _resolve_scope(scope_kb_id: str | None, scope_thread_id: str | None) -> str | None:
def _resolve_scope(
scope_kb_id: str | None,
scope_thread_id: str | None,
scope_project_id: str | None = None,
) -> str | list[str] | None:
"""KB (an explicit pick) is exclusive; project and thread scopes combine so a
project chat also retrieves from its own attached documents."""
if scope_kb_id:
return kb_scope(scope_kb_id)
scopes = []
if scope_project_id:
scopes.append(project_scope(scope_project_id))
if scope_thread_id:
return thread_scope(scope_thread_id)
return None
scopes.append(thread_scope(scope_thread_id))
if not scopes:
return None
return scopes[0] if len(scopes) == 1 else scopes
def _format(rows, hits) -> tuple[str, list[dict]]:
@ -83,6 +95,7 @@ def search_knowledge_base_with_sources(
query: str,
scope_kb_id: str | None = None,
scope_thread_id: str | None = None,
scope_project_id: str | None = None,
top_k: int | None = None,
min_score: float = 0.0,
model_name: str | None = None,
@ -92,7 +105,7 @@ def search_knowledge_base_with_sources(
rendered ``<chunk>`` block's ``id``."""
if not query or not query.strip():
return "Error: query is empty.", []
scope = _resolve_scope(scope_kb_id, scope_thread_id)
scope = _resolve_scope(scope_kb_id, scope_thread_id, scope_project_id)
if scope is None:
return "No documents are attached to this chat.", []
@ -124,6 +137,7 @@ def search_for_autoinject(
query: str,
scope_kb_id: str | None = None,
scope_thread_id: str | None = None,
scope_project_id: str | None = None,
top_k: int | None = None,
min_dense_score: float = 0.70,
model_name: str | None = None,
@ -138,7 +152,7 @@ def search_for_autoinject(
"""
if not query or not query.strip():
return None
scope = _resolve_scope(scope_kb_id, scope_thread_id)
scope = _resolve_scope(scope_kb_id, scope_thread_id, scope_project_id)
if scope is None:
return None
k = top_k or config.TOP_K_HYBRID
@ -177,6 +191,7 @@ def search_knowledge_base(
query: str,
scope_kb_id: str | None = None,
scope_thread_id: str | None = None,
scope_project_id: str | None = None,
top_k: int | None = None,
min_score: float = 0.0,
model_name: str | None = None,
@ -186,6 +201,7 @@ def search_knowledge_base(
query = query,
scope_kb_id = scope_kb_id,
scope_thread_id = scope_thread_id,
scope_project_id = scope_project_id,
top_k = top_k,
min_score = min_score,
model_name = model_name,

21
studio/backend/core/training/resume.py
View file

@ -3,6 +3,7 @@
"""Helpers for validating resumable training outputs."""
import json
from pathlib import Path
from typing import Optional
@ -56,9 +57,29 @@ def normalize_resume_output_dir(path_value: str) -> str:
return str(path)
def _run_config(run: dict) -> dict:
raw_config = run.get("config_json")
if isinstance(raw_config, dict):
return raw_config
if not isinstance(raw_config, str) or not raw_config.strip():
return {}
try:
parsed = json.loads(raw_config)
except (json.JSONDecodeError, TypeError):
return {}
return parsed if isinstance(parsed, dict) else {}
def _uses_s3_dataset(run: dict) -> bool:
config = _run_config(run)
return config.get("dataset_source") == "s3" or "s3_dataset" in config
def can_resume_run(run: dict) -> bool:
if run.get("resumed_later"):
return False
if _uses_s3_dataset(run):
return False
final_step = run.get("final_step")
total_steps = run.get("total_steps")

228
studio/backend/core/training/s3_dataset.py Normal file
View file

@ -0,0 +1,228 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
S3 dataset loader.
Downloads dataset files (parquet / json / jsonl / csv) from an AWS S3 bucket
to a local temp directory so the existing local-file dataset path can consume
them. boto3 is an optional dependency and is imported lazily — callers should
gate on :func:`boto3_available` before invoking the loader.
The S3 config dict mirrors ``models.training.S3Config.model_dump()`` (snake_case
keys): bucket, region, prefix, access_key_id, secret_access_key, use_iam_role.
Credentials are read once to build the client and never logged or persisted.
"""
from __future__ import annotations
import logging
import os
import shutil
import tempfile
from importlib.util import find_spec
from typing import Callable, Optional
logger = logging.getLogger(__name__)
# Extensions the local-file loader (UnslothTrainer._loader_for_files) understands.
SUPPORTED_EXTENSIONS = (".parquet", ".json", ".jsonl", ".csv")
_JSON_EXTENSIONS = (".json", ".jsonl")
_IGNORED_METADATA_FILENAMES = {
"dataset_info.json",
"metadata.json",
"schema.json",
"state.json",
}
class S3DownloadCancelled(RuntimeError):
"""Raised when the caller cancels an S3 dataset download."""
class S3DatasetDownload:
def __init__(
self,
files: list[str],
temp_dir: Optional[str] = None,
):
self.files = files
self.temp_dir = temp_dir
def cleanup(self) -> None:
if not self.temp_dir:
return
shutil.rmtree(self.temp_dir, ignore_errors = True)
self.temp_dir = None
def boto3_available() -> bool:
"""True if boto3 can be imported (without importing it)."""
return find_spec("boto3") is not None
def _build_s3_client(s3_config: dict):
"""Create a boto3 S3 client from the config dict.
Uses explicit access keys when provided, otherwise falls back to the
default credential chain (IAM role / instance profile / env / shared creds).
"""
import boto3 # lazy: optional dependency
region = s3_config.get("region") or "us-east-1"
use_iam_role = bool(s3_config.get("use_iam_role"))
access_key_id = s3_config.get("access_key_id")
secret_access_key = s3_config.get("secret_access_key")
if not use_iam_role and access_key_id and secret_access_key:
return boto3.client(
"s3",
region_name = region,
aws_access_key_id = access_key_id,
aws_secret_access_key = secret_access_key,
)
# IAM role / instance profile / ambient credentials
return boto3.client("s3", region_name = region)
def _list_dataset_keys(client, bucket: str, prefix: Optional[str]) -> list[str]:
"""List object keys under ``prefix`` that have a supported data extension."""
paginator = client.get_paginator("list_objects_v2")
list_kwargs = {"Bucket": bucket}
if prefix:
list_kwargs["Prefix"] = prefix
keys: list[str] = []
for page in paginator.paginate(**list_kwargs):
for obj in page.get("Contents", []):
key = obj["Key"]
if key.endswith("/"):
continue # directory placeholder
if os.path.basename(key).lower() in _IGNORED_METADATA_FILENAMES:
continue
if key.lower().endswith(SUPPORTED_EXTENSIONS):
keys.append(key)
return keys
def _extension_family(key: str) -> str:
ext = os.path.splitext(key)[1].lower()
if ext in _JSON_EXTENSIONS:
return "json"
return ext.lstrip(".")
def _validate_single_extension_family(keys: list[str]) -> None:
families: list[str] = []
for key in keys:
family = _extension_family(key)
if family not in families:
families.append(family)
if len(families) <= 1:
return
raise ValueError(
"S3 prefix contains mixed dataset formats "
f"({', '.join(families)}). Keep one dataset format under the selected prefix."
)
def _unique_local_path(target_dir: str, filename: str, used_paths: set[str]) -> str:
"""Return an unused flattened path for an S3 object basename."""
stem, ext = os.path.splitext(filename)
candidate = os.path.join(target_dir, filename)
suffix = 1
while candidate in used_paths or os.path.exists(candidate):
candidate = os.path.join(target_dir, f"{stem}_{suffix}{ext}")
suffix += 1
used_paths.add(candidate)
return candidate
def _raise_if_cancelled(cancel_callback: Optional[Callable[[], bool]]) -> None:
if cancel_callback is not None and cancel_callback():
raise S3DownloadCancelled("S3 dataset download cancelled")
def prepare_s3_dataset_download(
s3_config: dict,
dest_dir: Optional[str] = None,
cancel_callback: Optional[Callable[[], bool]] = None,
) -> S3DatasetDownload:
"""Download supported dataset files from S3 to a local directory.
Returns the local files plus the owned temporary directory, when one was
created. Call ``cleanup()`` after the dataset loader has materialized data.
Raises ``RuntimeError`` if boto3 is missing, and ``ValueError`` if the
bucket/prefix contains no supported dataset files.
"""
if not boto3_available():
raise RuntimeError("S3 dataset loading requires boto3. Install it with: pip install boto3")
bucket = s3_config.get("bucket")
if not bucket:
raise ValueError("s3_config.bucket is required")
prefix = s3_config.get("prefix")
_raise_if_cancelled(cancel_callback)
client = _build_s3_client(s3_config)
keys = _list_dataset_keys(client, bucket, prefix)
_raise_if_cancelled(cancel_callback)
if not keys:
where = f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}"
raise ValueError(
f"No supported dataset files ({', '.join(SUPPORTED_EXTENSIONS)}) "
f"found under {where}"
)
_validate_single_extension_family(keys)
owns_temp_dir = dest_dir is None
target_dir = dest_dir or tempfile.mkdtemp(prefix = "unsloth_s3_dataset_")
try:
os.makedirs(target_dir, exist_ok = True)
local_files: list[str] = []
used_paths: set[str] = set()
for key in keys:
_raise_if_cancelled(cancel_callback)
filename = os.path.basename(key)
local_path = _unique_local_path(target_dir, filename, used_paths)
download_kwargs = {}
if cancel_callback is not None:
download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(cancel_callback)
client.download_file(bucket, key, local_path, **download_kwargs)
_raise_if_cancelled(cancel_callback)
local_files.append(local_path)
except Exception:
if owns_temp_dir:
shutil.rmtree(target_dir, ignore_errors = True)
raise
logger.info(
"Downloaded %d dataset file(s) from s3://%s/%s to %s",
len(local_files),
bucket,
prefix or "",
target_dir,
)
return S3DatasetDownload(
files = local_files,
temp_dir = target_dir if owns_temp_dir else None,
)
def download_s3_dataset(
s3_config: dict,
dest_dir: Optional[str] = None,
cancel_callback: Optional[Callable[[], bool]] = None,
) -> list[str]:
download = prepare_s3_dataset_download(
s3_config,
dest_dir = dest_dir,
cancel_callback = cancel_callback,
)
return download.files

44
studio/backend/core/training/trainer.py
View file

@ -2227,6 +2227,7 @@ class UnslothTrainer:
dataset_slice_start: int = None,
dataset_slice_end: int = None,
is_cpt: bool = False,
s3_config: dict = None,
) -> Optional[tuple]:
"""
Load and prepare a dataset for training.
@ -2237,6 +2238,9 @@ class UnslothTrainer:
Returns (dataset_info, eval_dataset) or None on error; eval_dataset
may be None if no eval split is available.
"""
from core.training.s3_dataset import S3DownloadCancelled
s3_download = None
try:
dataset = None
eval_dataset = None
@ -2272,6 +2276,22 @@ class UnslothTrainer:
return result.dataset
# S3 datasets are downloaded to a local temp dir and then consumed
# through the same local-file path below.
if s3_config and not local_datasets:
from core.training.s3_dataset import prepare_s3_dataset_download
self._update_progress(status_message = "Downloading dataset from S3...")
s3_download = prepare_s3_dataset_download(
s3_config,
cancel_callback = lambda: self.should_stop,
)
local_datasets = s3_download.files
if self.should_stop:
logger.info("Stopped during S3 download\n")
return None
logger.info(f"Downloaded {len(local_datasets)} file(s) from S3\n")
if local_datasets:
# Use load_dataset() for an Arrow-backed result; in-memory
# Dataset.from_list() has no cache and forces num_proc=1 during
@ -2539,10 +2559,16 @@ class UnslothTrainer:
return (dataset_info, eval_dataset)
except S3DownloadCancelled:
logger.info("Stopped during S3 download\n")
return None
except Exception as e:
logger.error(f"Error loading dataset: {e}")
self._update_progress(error = str(e))
return None
finally:
if s3_download is not None:
s3_download.cleanup()
def _auto_detect_eval_split_from_hf(
self, dataset_source: str, subset: str
@ -3367,7 +3393,19 @@ class UnslothTrainer:
# ========== PROGRESS TRACKING ==========
self.trainer.add_callback(self._create_progress_callback())
num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset)
num_samples = None
if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None:
try:
num_samples = len(self.trainer.train_dataset)
except TypeError:
logger.debug(
"train_dataset does not support len(); falling back to "
"raw dataset size for step estimation."
)
if num_samples is None:
num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset)
batch_size = training_args.get("batch_size", 2)
total_steps = self._calculate_total_steps(
num_samples,
@ -3376,10 +3414,8 @@ class UnslothTrainer:
training_args.get("num_epochs", 3),
training_args.get("max_steps", 0),
)
self._update_progress(total_steps = total_steps)
# ========== START TRAINING ==========
self._update_progress(status_message = "Starting training...")
self._update_progress(total_steps = total_steps, status_message = "Starting training...")
logger.info("Starting training...\n")
self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint"))

86
studio/backend/core/training/training.py
View file

@ -37,9 +37,73 @@ from utils.paths import outputs_root
logger = get_logger(__name__)
def _coerce_seed(value, default = 3407) -> int:
"""Normalize None / non-int to `default` (transformers.set_seed(None) raises)."""
if value is None:
return int(default)
try:
return int(value)
except (TypeError, ValueError):
return int(default)
def _coerce_optional_bool(value, default: bool) -> bool:
"""Treat explicit None as `default` instead of `bool(None) == False`."""
if value is None:
return bool(default)
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in ("true", "1", "yes", "on"):
return True
if normalized in ("false", "0", "no", "off", ""):
return False
return bool(value)
def _coerce_optional_nonneg_float(name: str, value):
"""Reject negatives; HTTP `ge=0` doesn't cover raw `**kwargs` callers."""
if value is None:
return None
try:
coerced = float(value)
except (TypeError, ValueError):
raise ValueError(f"Unsloth: {name}={value!r} must be a non-negative float or None.")
if coerced < 0:
raise ValueError(f"Unsloth: {name}={coerced} must be >= 0 (use 0 or None to disable).")
return coerced
_HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$")
def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]:
db_config = {
k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"}
}
s3_config = config.get("s3_config")
if hasattr(s3_config, "model_dump"):
s3_config = s3_config.model_dump()
if isinstance(s3_config, dict) and s3_config:
db_config["dataset_source"] = "s3"
db_config["s3_dataset"] = {
"bucket": s3_config.get("bucket"),
"region": s3_config.get("region"),
"prefix": s3_config.get("prefix"),
"use_iam_role": bool(s3_config.get("use_iam_role")),
}
return db_config
def _s3_dataset_name(s3_dataset: Any) -> Optional[str]:
if not isinstance(s3_dataset, dict):
return None
bucket = s3_dataset.get("bucket")
if not bucket:
return None
prefix = s3_dataset.get("prefix")
return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}"
def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
"""Remove only HF Trainer ``tmp-checkpoint-<step>/`` partials after a cancel.
@ -211,7 +275,17 @@ class TrainingBackend:
"save_steps": kwargs.get("save_steps", 0),
"weight_decay": kwargs.get("weight_decay", 0.001),
"max_grad_norm": kwargs.get("max_grad_norm", 0.0),
"random_seed": kwargs.get("random_seed", 3407),
"max_grad_value": _coerce_optional_nonneg_float(
"max_grad_value", kwargs.get("max_grad_value")
),
"max_grad_leaf_norm": _coerce_optional_nonneg_float(
"max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm")
),
"cast_norm_output_to_input_dtype": _coerce_optional_bool(
kwargs.get("cast_norm_output_to_input_dtype"), True
),
# MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises).
"random_seed": _coerce_seed(kwargs.get("random_seed")),
"packing": kwargs.get("packing", False),
"optim": kwargs.get("optim", "adamw_8bit"),
"lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"),
@ -236,6 +310,7 @@ class TrainingBackend:
"resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
"trust_remote_code": kwargs.get("trust_remote_code", False),
"gpu_ids": kwargs.get("gpu_ids"),
"s3_config": kwargs.get("s3_config"),
}
# Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request.
@ -309,7 +384,7 @@ class TrainingBackend:
self._run_finalized = False
self._db_run_created = False
self._db_total_steps_set = False
self._db_config = {k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}}
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Assign subprocess handles after state reset.
@ -732,8 +807,11 @@ class TrainingBackend:
try:
from storage.studio_db import create_run
dataset_name = self._db_config.get("hf_dataset") or next(
iter(self._db_config.get("local_datasets") or []), "unknown"
dataset_name = (
self._db_config.get("hf_dataset")
or next(iter(self._db_config.get("local_datasets") or []), None)
or _s3_dataset_name(self._db_config.get("s3_dataset"))
or "unknown"
)
create_run(
id = self.current_job_id,

269
studio/backend/core/training/worker.py
View file

@ -1362,6 +1362,32 @@ def _run_mlx_training(event_queue, stop_queue, config):
kwargs["message"] = sm
event_queue.put({"type": event_type, "ts": time.time(), **kwargs})
_stop_save = [True]
_stop_requested = [False]
_trainer_ref = [None]
def _is_stop_requested():
return _stop_requested[0]
def _poll_stop():
while True:
try:
msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
_stop_save[0] = msg.get("save", True)
_stop_requested[0] = True
trainer = _trainer_ref[0]
if trainer is not None:
trainer.stop_requested = True
return
except _queue.Empty:
continue
except (EOFError, OSError):
return
stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
_send("status", status_message = "Loading MLX libraries...")
import mlx.core as mx
@ -1424,6 +1450,14 @@ def _run_mlx_training(event_queue, stop_queue, config):
is_dataset_image = bool(config.get("is_dataset_image", False))
training_type = config.get("training_type", "LoRA/QLoRA")
use_lora = training_type == "LoRA/QLoRA"
# Normalize seed; explicit None must not reach the seed chain.
_raw_seed = config.get("random_seed", 3407)
random_seed = 3407 if _raw_seed is None else int(_raw_seed)
# `config.get(k, d)` only fills d when key is missing; handle explicit None too.
_model_seed = config.get("model_random_state")
model_random_state = random_seed if _model_seed is None else int(_model_seed)
_lora_seed = config.get("lora_random_state")
lora_random_state = random_seed if _lora_seed is None else int(_lora_seed)
model, tokenizer = FastMLXModel.from_pretrained(
model_name,
load_in_4bit = config.get("load_in_4bit", True),
@ -1431,7 +1465,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
text_only = None if is_dataset_image else True,
token = hf_token,
trust_remote_code = bool(config.get("trust_remote_code", False)),
random_state = config.get("random_seed", 3407),
random_state = model_random_state,
)
is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False))
@ -1473,7 +1507,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
lora_dropout = config.get("lora_dropout", 0.0),
use_rslora = config.get("use_rslora", False),
init_lora_weights = config.get("init_lora_weights", True),
random_state = config.get("random_seed", 3407),
random_state = lora_random_state,
target_modules = config.get("target_modules")
or [
"q_proj",
@ -1541,6 +1575,26 @@ def _run_mlx_training(event_queue, stop_queue, config):
elif config.get("local_datasets"):
dataset = _load_local(config["local_datasets"])
dataset = _slice(dataset)
elif config.get("s3_config"):
from core.training.s3_dataset import (
S3DownloadCancelled,
prepare_s3_dataset_download,
)
_send("status", status_message = "Downloading dataset from S3...")
try:
s3_download = prepare_s3_dataset_download(
config["s3_config"],
cancel_callback = _is_stop_requested,
)
try:
dataset = _load_local(s3_download.files)
finally:
s3_download.cleanup()
except S3DownloadCancelled:
_send("complete", output_dir = None, status_message = "Training cancelled")
return
dataset = _slice(dataset)
else:
raise ValueError("No dataset specified")
@ -1667,12 +1721,12 @@ def _run_mlx_training(event_queue, stop_queue, config):
warmup_steps = 5
# ── 5. Build output dir ──
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import resolve_output_dir, ensure_dir, default_run_dir_name
output_dir = config.get("output_dir", "")
if not output_dir:
output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}"
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import resolve_output_dir, ensure_dir
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
@ -1684,41 +1738,80 @@ def _run_mlx_training(event_queue, stop_queue, config):
else:
eval_steps_val = int(eval_steps_val)
# MLX: per-element clip to [-1, 1]; norm clip disabled (its global reduction
# breaks MLX's eager pipeline). 1.0 not 5.0: |g_i| > 5 rarely fires, so the
# historical 5.0 was effectively a no-op.
# Per-element clipping only; trainer owns the None default. Re-validate
# for direct worker callers (training.py normalizes the main path).
max_grad_norm = 0.0
max_grad_value = 1.0 # TODO: expose MLX grad-clip in Studio UI for power users
max_grad_value = config.get("max_grad_value")
if max_grad_value is not None:
max_grad_value = float(max_grad_value)
if max_grad_value < 0:
raise ValueError(
f"Unsloth MLX: max_grad_value={max_grad_value} must be >= 0 "
"(0 or None disables elementwise clipping)."
)
max_grad_leaf_norm = config.get("max_grad_leaf_norm")
if max_grad_leaf_norm is not None:
max_grad_leaf_norm = float(max_grad_leaf_norm)
if max_grad_leaf_norm < 0:
raise ValueError(
f"Unsloth MLX: max_grad_leaf_norm={max_grad_leaf_norm} must be >= 0 "
"(0 or None disables proportional leaf-norm clipping)."
)
weight_decay = config.get("weight_decay", 0.001)
weight_decay = 0.001 if weight_decay is None else float(weight_decay)
mlx_config_kwargs = dict(
per_device_train_batch_size = batch_size,
gradient_accumulation_steps = grad_accum,
max_steps = max_steps,
learning_rate = lr_value,
warmup_steps = warmup_steps,
lr_scheduler_type = lr_scheduler_type,
optim = optim_name,
weight_decay = weight_decay,
max_grad_norm = max_grad_norm,
max_grad_value = max_grad_value,
logging_steps = 1,
max_seq_length = max_seq_length,
seed = random_seed,
use_cce = True,
compile = True,
gradient_checkpointing = use_grad_checkpoint,
streaming = is_vlm,
packing = bool(config.get("packing", False)),
output_dir = output_dir,
save_steps = int(config.get("save_steps", 0) or 0),
eval_steps = eval_steps_val,
)
# Feature-detect optional fields so this PR works without the paired zoo bump.
_supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {})
if "cast_norm_output_to_input_dtype" in _supported_fields:
# Explicit None falls back to True (default).
_raw_cast = config.get("cast_norm_output_to_input_dtype", True)
mlx_config_kwargs["cast_norm_output_to_input_dtype"] = (
True if _raw_cast is None else bool(_raw_cast)
)
if "dataset_order" in _supported_fields:
mlx_config_kwargs["dataset_order"] = "torch_randperm"
if "max_grad_leaf_norm" in _supported_fields:
mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm
if "append_eos" in _supported_fields:
raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw"
# Studio SFT formatting owns rendered examples; raw/CPT text still
# needs MLX to append EOS like the CUDA raw-text path.
mlx_config_kwargs["append_eos"] = bool(raw_text_mode)
trainer = MLXTrainer(
model = model,
tokenizer = tokenizer,
train_dataset = dataset,
eval_dataset = eval_dataset,
args = MLXTrainingConfig(
per_device_train_batch_size = batch_size,
gradient_accumulation_steps = grad_accum,
max_steps = max_steps,
learning_rate = lr_value,
warmup_steps = warmup_steps,
lr_scheduler_type = lr_scheduler_type,
optim = optim_name,
weight_decay = float(config.get("weight_decay", 0.001) or 0.001),
max_grad_norm = max_grad_norm,
max_grad_value = max_grad_value,
logging_steps = 1,
max_seq_length = max_seq_length,
seed = config.get("random_seed", 3407),
use_cce = True,
compile = True,
gradient_checkpointing = use_grad_checkpoint,
streaming = is_vlm,
packing = bool(config.get("packing", False)),
output_dir = output_dir,
save_steps = int(config.get("save_steps", 0) or 0),
eval_steps = eval_steps_val,
),
args = MLXTrainingConfig(**mlx_config_kwargs),
)
_trainer_ref[0] = trainer
if _stop_requested[0]:
trainer.stop_requested = True
# Tell the parent eval is configured so the frontend shows the eval chart
if eval_dataset is not None and eval_steps_val > 0:
@ -1759,7 +1852,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
wandb_token = config.get("wandb_token")
if wandb_token:
os.environ["WANDB_API_KEY"] = wandb_token
_wandb_sensitive = {"hf_token", "wandb_token"}
_wandb_sensitive = {"hf_token", "wandb_token", "s3_config"}
wandb_run = _wandb.init(
project = config.get("wandb_project") or "unsloth-mlx",
config = {k: v for k, v in config.items() if k not in _wandb_sensitive},
@ -1861,26 +1954,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.add_eval_callback(_on_eval)
# ── 10. Stop signal polling ──
_stop_save = [True] # mutable so thread can update; [save_flag]
def _poll_stop():
while True:
try:
msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
_stop_save[0] = msg.get("save", True)
trainer.stop_requested = True
return
except _queue.Empty:
continue
except (EOFError, OSError):
# Safe: pipe permanently broken, no more messages can arrive.
return
stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
# ── 11. Run training ──
gc.collect()
mx.synchronize()
@ -2515,6 +2588,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
resolve_output_dir,
resolve_tensorboard_dir,
datasets_root,
default_run_dir_name,
)
import transformers
@ -2556,7 +2630,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
def _on_progress(progress: TrainingProgress):
has_train_loss = progress.step > 0 and progress.loss is not None
has_eval_loss = progress.eval_loss is not None
if has_train_loss or has_eval_loss:
if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss:
event_queue.put(
{
"type": "progress",
@ -2640,6 +2714,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
dataset_slice_start = config.get("dataset_slice_start"),
dataset_slice_end = config.get("dataset_slice_end"),
is_cpt = _is_cpt_for_dataset,
s3_config = config.get("s3_config"),
)
if isinstance(dataset_result, tuple):
@ -2837,7 +2912,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}"
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
@ -2988,7 +3063,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
from datasets import Dataset
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
from transformers import TrainerCallback
from utils.paths import datasets_root, resolve_output_dir
from utils.paths import datasets_root, resolve_output_dir, default_run_dir_name
except ImportError as e:
event_queue.put(
{
@ -3104,20 +3179,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
subset = config.get("subset") or None
train_split = config.get("train_split", "train") or "train"
if hf_dataset and hf_dataset.strip():
hf_token = config.get("hf_token", "")
hf_token = hf_token if hf_token and hf_token.strip() else None
dataset = load_dataset(
hf_dataset.strip(),
subset,
split = train_split,
token = hf_token,
)
elif local_datasets:
# Load local file(s) — mirrors the non-embedding pipeline's directory
# handling so recipe outputs (parquet-files/) work.
def _load_local_embedding_dataset(dataset_paths: list[str]):
all_files: list[str] = []
for dataset_file in local_datasets:
for dataset_file in dataset_paths:
file_path = (
dataset_file
if os.path.isabs(dataset_file)
@ -3147,17 +3211,58 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
else:
all_files.append(file_path)
if all_files:
first_ext = Path(all_files[0]).suffix.lower()
if first_ext in (".json", ".jsonl"):
loader = "json"
elif first_ext == ".csv":
loader = "csv"
elif first_ext == ".parquet":
loader = "parquet"
else:
raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
dataset = load_dataset(loader, data_files = all_files, split = "train")
if not all_files:
raise ValueError("No local dataset files found")
first_ext = Path(all_files[0]).suffix.lower()
if first_ext in (".json", ".jsonl"):
loader = "json"
elif first_ext == ".csv":
loader = "csv"
elif first_ext == ".parquet":
loader = "parquet"
else:
raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
return load_dataset(loader, data_files = all_files, split = "train")
if hf_dataset and hf_dataset.strip():
hf_token = config.get("hf_token", "")
hf_token = hf_token if hf_token and hf_token.strip() else None
dataset = load_dataset(
hf_dataset.strip(),
subset,
split = train_split,
token = hf_token,
)
elif local_datasets:
dataset = _load_local_embedding_dataset(local_datasets)
elif config.get("s3_config"):
from core.training.s3_dataset import (
S3DownloadCancelled,
prepare_s3_dataset_download,
)
_send_status(event_queue, "Downloading dataset from S3...")
s3_download = None
try:
s3_download = prepare_s3_dataset_download(
config["s3_config"],
cancel_callback = lambda: _should_stop,
)
dataset = _load_local_embedding_dataset(s3_download.files)
except S3DownloadCancelled:
event_queue.put(
{
"type": "complete",
"output_dir": None,
"status_message": "Training cancelled",
"ts": time.time(),
}
)
return
finally:
if s3_download is not None:
s3_download.cleanup()
else:
event_queue.put(
{
@ -3216,7 +3321,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = str(resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}"))
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = str(resolve_output_dir(output_dir))
num_epochs = config.get("num_epochs", 2)

19
studio/backend/hub/utils/inventory_scan.py
View file

@ -142,13 +142,18 @@ def _compute_all_hf_cache_scans() -> list:
logger.warning("Could not scan active HF cache: %s", exc)
for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir):
extra = extra_fn()
if extra.is_dir() and str(extra.resolve()) not in seen:
seen.add(str(extra.resolve()))
try:
scans.append(scan_cache_dir(cache_dir = str(extra)))
except Exception as exc:
logger.warning("Could not scan HF cache %s: %s", extra, exc)
try:
extra = extra_fn()
# is_dir()/resolve() can raise on an inaccessible path; skip it.
if not extra.is_dir():
continue
resolved = str(extra.resolve())
if resolved in seen:
continue
seen.add(resolved)
scans.append(scan_cache_dir(cache_dir = str(extra)))
except Exception as exc:
logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc)
return scans

5
studio/backend/main.py
View file

@ -193,6 +193,11 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth`
# does) so its lazy submodule imports (export, hardware, mlx) and the
# DiffusionGemma runner never trip the install guard on a clean install.
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
import hashlib
import mimetypes
import re as _re

20
studio/backend/models/export.py
View file

@ -3,14 +3,14 @@
"""Pydantic schemas for Export API."""
from pathlib import Path
from pathlib import Path, PureWindowsPath
from pydantic import BaseModel, Field, field_validator
from typing import List, Optional, Literal, Dict, Any
def _validate_save_directory(value: str) -> str:
"""Reject save_directory values that escape the export root."""
"""Validate save_directory — allows absolute paths (user may want a different drive)."""
if value is None:
raise ValueError("save_directory is required")
raw = str(value).strip()
@ -20,15 +20,15 @@ def _validate_save_directory(value: str) -> str:
raise ValueError("save_directory may not contain null bytes")
if any(ch in raw for ch in ("\r", "\n")):
raise ValueError("save_directory may not contain control characters")
if len(raw) > 255:
raise ValueError("save_directory must be <= 255 characters")
path = Path(raw).expanduser()
if path.is_absolute():
raise ValueError(
"save_directory must be a name or relative path under the "
"export root; absolute paths are rejected"
)
if ".." in path.parts:
path_parts = (*path.parts, *PureWindowsPath(raw).parts, *raw.replace("\\", "/").split("/"))
if any(len(part) > 255 for part in path_parts if part not in ("", ".", "/", "\\")):
raise ValueError("save_directory path components must be <= 255 characters")
if (
".." in path.parts
or ".." in PureWindowsPath(raw).parts
or ".." in raw.replace("\\", "/").split("/")
):
raise ValueError("save_directory may not contain '..' segments")
return raw

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

@ -87,6 +87,15 @@ class LoadRequest(BaseModel):
"'mtp' or 'mtp+ngram'."
),
)
tensor_parallel: bool = Field(
False,
description = (
"Split the model across GPUs by tensor (--split-mode tensor) "
"instead of by layer for GGUF models. Only affects multi-GPU "
"setups, where it can make generation significantly faster. "
"No effect on a single GPU. Ignored for non-GGUF models."
),
)
llama_extra_args: Optional[List[str]] = Field(
None,
description = (
@ -160,6 +169,9 @@ class LoadResponse(BaseModel):
is_vision: bool = Field(False, description = "Whether model is a vision model")
is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)")
is_diffusion: bool = Field(
False, description = "Whether model is a block-diffusion model (DiffusionGemma)"
)
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)")
@ -224,6 +236,10 @@ class LoadResponse(BaseModel):
"None when the platform default is in effect."
),
)
tensor_parallel: bool = Field(
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
class UnloadResponse(BaseModel):
@ -273,6 +289,9 @@ class InferenceStatusResponse(BaseModel):
)
is_vision: bool = Field(False, description = "Whether the active model is a vision model")
is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)")
is_diffusion: bool = Field(
False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)"
)
gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)")
is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model")
audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac")
@ -339,6 +358,10 @@ class InferenceStatusResponse(BaseModel):
"None when the platform default is in effect."
),
)
tensor_parallel: bool = Field(
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
llama_cpp_supports_mtp: bool = Field(
True,
description = (
@ -393,7 +416,7 @@ class ImageUrl(BaseModel):
"""Image URL object — supports data URIs and remote URLs."""
url: str = Field(..., description = "data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
detail: Optional[Literal["auto", "low", "high", "original"]] = "auto"
class ImageContentPart(BaseModel):
@ -690,6 +713,10 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.",
)
confirm_tool_calls: Optional[bool] = Field(
None,
description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@ -926,6 +953,12 @@ class ChatCompletionRequest(BaseModel):
return self
class ToolConfirmRequest(BaseModel):
session_id: Optional[str] = None
approval_id: Optional[str] = None
decision: Literal["allow", "deny"] = "deny"
# ── OpenAI shell-tool container management ─────────────────────
@ -1092,7 +1125,7 @@ class ResponsesInputImagePart(BaseModel):
type: Literal["input_image"]
image_url: str = Field(..., description = "data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
detail: Optional[Literal["auto", "low", "high", "original"]] = "auto"
class ResponsesOutputTextPart(BaseModel):
@ -1304,6 +1337,23 @@ class ResponsesOutputMessage(BaseModel):
content: list[ResponsesOutputTextContent] = Field(default_factory = list)
class ResponsesOutputReasoningContent(BaseModel):
"""A reasoning text content block inside a reasoning output item."""
type: Literal["reasoning_text"] = "reasoning_text"
text: str
class ResponsesOutputReasoning(BaseModel):
"""A top-level reasoning output item in the Responses API response."""
type: Literal["reasoning"] = "reasoning"
id: str = Field(default_factory = lambda: f"rs_{uuid.uuid4().hex[:12]}")
status: Literal["completed", "in_progress", "incomplete"] = "completed"
summary: list = Field(default_factory = list)
content: Optional[list[ResponsesOutputReasoningContent]] = None
class ResponsesOutputFunctionCall(BaseModel):
"""A function-call output item in the Responses API response.
@ -1318,7 +1368,11 @@ class ResponsesOutputFunctionCall(BaseModel):
status: Literal["completed", "in_progress", "incomplete"] = "completed"
ResponsesOutputItem = Union[ResponsesOutputMessage, ResponsesOutputFunctionCall]
ResponsesOutputItem = Union[
ResponsesOutputMessage,
ResponsesOutputReasoning,
ResponsesOutputFunctionCall,
]
class ResponsesUsage(BaseModel):

3
studio/backend/models/providers.py
View file

@ -108,6 +108,9 @@ class ProviderTestRequest(BaseModel):
base_url: Optional[str] = Field(
None, description = "Custom base URL (overrides registry default)"
)
model_id: Optional[str] = Field(
None, description = "Model ID for providers that need a chat probe"
)
class ProviderTestResult(BaseModel):

75
studio/backend/models/training.py
View file

@ -29,6 +29,43 @@ _MIN_VISION_IMAGE_SIZE = 256
_MAX_VISION_IMAGE_SIZE = 2048
class S3Config(BaseModel):
"""S3 bucket configuration for loading datasets from AWS S3"""
# Accept both snake_case and the frontend's camelCase field names.
model_config = ConfigDict(populate_by_name = True)
bucket: str = Field(..., description = "S3 bucket name")
region: str = Field("us-east-1", description = "AWS region")
prefix: Optional[str] = Field(None, description = "Optional path prefix within bucket")
access_key_id: Optional[str] = Field(
None,
alias = "accessKeyId",
description = "AWS access key ID (optional if using IAM role)",
)
secret_access_key: Optional[str] = Field(
None,
alias = "secretAccessKey",
description = "AWS secret access key (optional if using IAM role)",
)
use_iam_role: bool = Field(
False,
alias = "useIamRole",
description = "Use IAM role credentials instead of access keys",
)
@model_validator(mode = "after")
def _check_credentials(self) -> "S3Config":
# Require either IAM role auth or a full key pair so credentials are
# never half-configured.
if not self.use_iam_role and not (self.access_key_id and self.secret_access_key):
raise ValueError(
"s3_config requires either use_iam_role=True or both "
"access_key_id and secret_access_key"
)
return self
def _parse_lr(v: Any) -> float:
"""Parse learning_rate as a positive float strictly below _MAX_LR_VALUE."""
if v is None:
@ -288,7 +325,37 @@ class TrainingStartRequest(BaseModel):
ge = 0,
description = "Global gradient norm clipping threshold. Set 0 to disable.",
)
random_seed: int = Field(42, description = "Random seed")
max_grad_value: Optional[float] = Field(
None,
ge = 0,
description = (
"MLX-only elementwise gradient value clipping threshold. "
"If unset, MLX uses its runtime default."
),
)
max_grad_leaf_norm: Optional[float] = Field(
None,
ge = 0,
description = (
"MLX-only proportional per-parameter gradient norm cap. "
"Preserves each tensor's gradient direction without global norm "
"clipping's memory overhead."
),
)
cast_norm_output_to_input_dtype: bool = Field(
True,
description = (
"MLX-only: keep norm parameters in fp32 but cast norm outputs "
"back to the incoming activation dtype."
),
)
random_seed: int = Field(
3407,
description = (
"Random seed; matches the Studio backend / MLX worker default "
"and unsloth's historical recommended value."
),
)
packing: bool = Field(False, description = "Enable sequence packing")
optim: str = Field("adamw_8bit", description = "Optimizer")
lr_scheduler_type: str = Field("linear", description = "Learning rate scheduler type")
@ -338,6 +405,12 @@ class TrainingStartRequest(BaseModel):
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.",
)
# S3 dataset source configuration
s3_config: Optional[S3Config] = Field(
None,
description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.",
)
@model_validator(mode = "after")
def _check_steps_or_epochs(self) -> "TrainingStartRequest":
# Each accepts 0 as "use the other"; both 0 means nothing to train.

6
studio/backend/requirements/no-torch-runtime.txt
View file

@ -68,3 +68,9 @@ trl>=0.18.2,!=0.19.0,<=0.24.0
sentence-transformers
cut_cross_entropy
pillow
# RAG store + document parsing, mirroring studio.txt. Pinned here because
# this file installs --no-deps; without them Studio runs with RAG disabled.
sqlite-vec==0.1.9
pymupdf==1.27.2.3
python-docx==1.2.0

1
studio/backend/requirements/studio.txt
View file

@ -17,6 +17,7 @@ structlog>=24.1.0
diceware
ddgs
cryptography>=42.0.0
boto3>=1.34.0 # optional: S3 dataset loading
httpx>=0.27.0
fastmcp>=3.0.2
# RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in

29
studio/backend/routes/chat_history.py
View file

@ -332,6 +332,35 @@ async def delete_project(
status_code = 404,
detail = f"Project {project_id} not found",
)
# Best-effort: drop the project's RAG sources (lazy import keeps RAG optional).
try:
import os
from storage import rag_db
if rag_db.RAG_AVAILABLE:
from core.rag import store as rag_store
from utils.paths import rag_uploads_root
uploads = os.path.realpath(str(rag_uploads_root()))
conn = rag_db.get_connection()
try:
scope = rag_store.project_scope(project_id)
for doc in rag_store.list_documents(conn, scope):
full = rag_store.get_document(conn, doc["id"]) or {}
rag_store.delete_document(conn, doc["id"])
stored = full.get("stored_path")
# Also remove the uploaded file; confined to the uploads root.
if stored:
target = os.path.realpath(stored)
if (
os.path.isfile(target)
and os.path.commonpath([uploads, target]) == uploads
):
os.remove(target)
finally:
conn.close()
except Exception: # noqa: BLE001 - source cleanup must not block project deletion
logger.warning("failed to delete RAG sources for project %s", project_id, exc_info = True)
return ChatProject(**project)

30
studio/backend/routes/export.py
View file

@ -154,19 +154,41 @@ async def get_export_status(current_subject: str = Depends(get_current_subject))
)
def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]:
"""Best-effort registration so absolute exports show up in local scans."""
try:
from storage.studio_db import add_scan_folder
folder = add_scan_folder(str(path))
return True, str(folder.get("path") or path)
except Exception as exc:
logger.warning("Could not register export scan folder %s: %s", path, exc)
return False, None
def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
"""Return the export path relative to exports_root, hiding the install path."""
"""Return relative export paths, keeping external absolute paths visible."""
if not output_path:
return None
try:
from utils.paths.storage_roots import exports_root
path = Path(output_path)
# If it's outside exports_root, return the full absolute path
# so users can find their files on a different drive.
if path.is_absolute():
try:
path.resolve().relative_to(exports_root().resolve())
except ValueError:
registered, registered_path = _try_register_external_export(path)
return {
"output_path": str(path),
"scan_folder_registered": registered,
"scan_folder_path": registered_path,
}
rel = os.path.relpath(output_path, exports_root())
if rel.startswith(".."):
rel = os.path.basename(output_path)
return {"output_path": rel}
except Exception:
return {"output_path": os.path.basename(output_path)}
return {"output_path": output_path}
@router.post("/export/merged", response_model = ExportOperationResponse)

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

File diff suppressed because it is too large Load diff

25
studio/backend/routes/mcp_servers.py
View file

@ -10,12 +10,16 @@ from fastapi import APIRouter, Depends, HTTPException
from auth.authentication import get_current_subject
from core.inference.mcp_client import (
TOOL_CACHE_INVALIDATING_FIELDS,
cache_tools,
clear_oauth_tokens_async,
invalidate_tool_cache,
is_stdio,
list_tools_async,
parse_server_headers,
parse_stdio_command,
probe_timeout,
record_probe_failure,
stdio_mcp_enabled,
)
from core.inference.mcp_config_import import parse_mcp_config
@ -198,6 +202,11 @@ async def update_mcp_server(
):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.update_server(server_id, changes)
# A new endpoint/auth makes cached tools wrong and disabling makes them
# unreachable, so drop them and let the next send re-probe; a rename
# leaves them valid.
if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS:
invalidate_tool_cache(server_id)
return _row_to_response(mcp_servers_db.get_server(server_id))
@ -209,6 +218,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c
if old.get("use_oauth"):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.delete_server(server_id)
invalidate_tool_cache(server_id)
@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)
@ -238,8 +248,23 @@ async def refresh_mcp_server_tools(
error = str(exc),
exc_info = True,
)
current = mcp_servers_db.get_server(server_id)
if current is not None and not any(
current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS
):
# Start the cool-off so the next chat send doesn't immediately re-hang
# on this server's timeout. If the row changed while the probe was
# awaiting, the failure belongs to the old config and must not park
# the newly edited server.
record_probe_failure(server_id, use_oauth)
return McpServerProbeResult(ok = False, error = safe_curated_detail(exc))
# Warm the chat-path cache so the next send skips re-probing.
current = mcp_servers_db.get_server(server_id)
if current is not None and not any(
current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS
):
cache_tools(server_id, tools)
return McpServerProbeResult(ok = True, tool_count = len(tools))

182
studio/backend/routes/models.py
View file

@ -51,6 +51,14 @@ def _is_hidden_model(*values: str | None) -> bool:
return any(v and any(n in v.lower() for n in needles) for v in values)
def _safe_resolve(path: Path) -> Optional[str]:
"""resolve() to a string, or None when the path is inaccessible."""
try:
return str(path.resolve())
except OSError:
return None
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
@ -676,9 +684,9 @@ async def list_local_models(
# trusted Path objects are used for FS access; the user string is
# used for matching only, never for path construction.
allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir]
if legacy_hf.is_dir():
if _safe_is_dir(legacy_hf):
allowed_roots.append(legacy_hf)
if hf_default.is_dir():
if _safe_is_dir(hf_default):
allowed_roots.append(hf_default)
try:
from utils.paths import studio_root, outputs_root
@ -702,15 +710,20 @@ async def list_local_models(
try:
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve():
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if (
hf_default.is_dir()
and hf_default.resolve() != hf_cache_dir.resolve()
and hf_default.resolve() != legacy_hf.resolve()
_safe_is_dir(hf_default)
and default_real != hf_cache_real
and default_real != legacy_real
):
local_models += _scan_hf_cache(hf_default)
@ -2069,13 +2082,10 @@ async def get_gguf_variants(
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
# Which variants are fully downloaded in the HF cache. For split
# GGUFs ALL shards must be present, so sum cached bytes per variant
# vs. the expected total. Cache dir casing may differ from the
# canonical repo_id, so match case-insensitively.
cached_bytes_by_quant: dict[str, int] = {}
# Per-snapshot so a split GGUF's shards must all sit in one snapshot;
# mmproj adapters are excluded so they can't inflate a quant's bytes.
cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = []
try:
import re as _re
from huggingface_hub import constants as hf_constants
if not _is_valid_repo_id(repo_id):
@ -2088,21 +2098,31 @@ async def get_gguf_variants(
snapshots = entry / "snapshots"
if snapshots.is_dir():
for snap in snapshots.iterdir():
by_quant: dict[str, int] = {}
for f in _iter_gguf_paths(snap):
q = _extract_quant_label(f.name)
cached_bytes_by_quant[q] = (
cached_bytes_by_quant.get(q, 0) + f.stat().st_size
)
if _is_mmproj_filename(f.name):
continue
try:
size = f.stat().st_size
except OSError:
continue # broken symlink / unreadable: skip
q = _extract_quant_label(f.name).lower()
by_quant[q] = by_quant.get(q, 0) + size
if by_quant:
cached_bytes_by_quant_per_snapshot.append(by_quant)
break
except Exception:
pass
def _is_fully_downloaded(variant) -> bool:
cached = cached_bytes_by_quant.get(variant.quant, 0)
if cached == 0 or variant.size_bytes == 0:
if variant.size_bytes == 0:
return False
# Rounding tolerance (symlinks vs real sizes).
return cached >= variant.size_bytes * 0.99
# Complete within one snapshot (tolerance for symlink size jitter).
quant = variant.quant.lower()
return any(
by_quant.get(quant, 0) >= variant.size_bytes * 0.99
for by_quant in cached_bytes_by_quant_per_snapshot
)
return GgufVariantsResponse(
repo_id = repo_id,
@ -2157,16 +2177,26 @@ async def get_gguf_download_progress(
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
# Completed .gguf files for this variant in snapshots.
# Exclude mmproj so a vision adapter can't satisfy a same-label
# main variant (e.g. mmproj-F16 vs an F16 weight).
for f in _iter_gguf_paths(entry):
if _is_mmproj_filename(f.name):
continue
fname = f.name.lower().replace("-", "").replace("_", "")
if not variant_lower or variant_lower in fname:
downloaded_bytes += f.stat().st_size
try:
downloaded_bytes += f.stat().st_size
except OSError:
continue # broken symlink / unreadable: skip
# In-progress (.incomplete) downloads in blobs.
blobs_dir = entry / "blobs"
if blobs_dir.is_dir():
for f in blobs_dir.iterdir():
if f.is_file() and f.name.endswith(".incomplete"):
in_progress_bytes += f.stat().st_size
try:
in_progress_bytes += f.stat().st_size
except OSError:
continue
break
total_progress_bytes = downloaded_bytes + in_progress_bytes
@ -2300,11 +2330,22 @@ def _get_repo_size_cached(repo_id: str) -> int:
def _all_hf_cache_scans():
"""scan_cache_dir results for the active, legacy, and default HF caches."""
"""scan_cache_dir for the active, legacy, and default HF caches.
Each probe is isolated: an unreadable auxiliary cache (permission denied,
broken symlink, OS-redirected ~/.cache) is skipped, not fatal, so the
Downloaded list never blanks out and downloads never leak into Recommended.
"""
from huggingface_hub import scan_cache_dir
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir
scans = [scan_cache_dir()]
scans = []
# Guard the active cache too: degrade to "no downloads" instead of raising.
try:
scans.append(scan_cache_dir())
except Exception as exc:
logger.warning("Could not scan active HF cache: %s", exc)
seen: set[str] = set()
try:
# Resolve the active cache dir for dedup.
@ -2314,13 +2355,18 @@ def _all_hf_cache_scans():
pass
for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir):
extra = extra_fn()
if extra.is_dir() and str(extra.resolve()) not in seen:
seen.add(str(extra.resolve()))
try:
scans.append(scan_cache_dir(cache_dir = str(extra)))
except Exception as exc:
logger.warning("Could not scan HF cache %s: %s", extra, exc)
try:
extra = extra_fn()
# is_dir()/resolve() can raise on an inaccessible path; skip it.
if not extra.is_dir():
continue
resolved = str(extra.resolve())
if resolved in seen:
continue
seen.add(resolved)
scans.append(scan_cache_dir(cache_dir = str(extra)))
except Exception as exc:
logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc)
return scans
@ -2379,6 +2425,38 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
def _blob_mtime(f) -> float:
"""Blob modification time in epoch seconds (0.0 if unknown).
Prefers HF metadata ``blob_last_modified``, falls back to stat(); uses
only mtimes (portable across Windows, macOS, Linux), never path parsing.
"""
ts = getattr(f, "blob_last_modified", None)
if isinstance(ts, (int, float)) and ts > 0:
return float(ts)
blob_path = getattr(f, "blob_path", None)
if blob_path:
try:
return float(Path(blob_path).stat().st_mtime)
except OSError:
pass
return 0.0
def _repo_gguf_last_modified(repo_info) -> float:
"""Newest mtime among a repo's primary (non-mmproj) GGUF blobs.
Drives the Downloaded list's "last downloaded" ordering and groups a
multi-quant repo by its most recently downloaded quant.
"""
latest = 0.0
for revision in repo_info.revisions:
for f in revision.files:
if _is_main_gguf_filename(f.file_name):
latest = max(latest, _blob_mtime(f))
return latest
@router.get("/cached-gguf")
async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
"""List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
@ -2399,17 +2477,30 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
continue
key = repo_id.lower()
existing = seen_lower.get(key)
last_modified = _repo_gguf_last_modified(repo_info)
if existing is None or total_size > existing["size_bytes"]:
seen_lower[key] = {
row = {
"repo_id": repo_id,
"size_bytes": total_size,
"cache_path": str(repo_info.repo_path),
}
# Keep the newest timestamp across duplicate caches;
# attach only when known so absent rows sort as oldest.
lm = max(last_modified, (existing or {}).get("last_modified", 0.0))
if lm > 0:
row["last_modified"] = lm
seen_lower[key] = row
elif last_modified > existing.get("last_modified", 0.0):
existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}")
continue
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
# Newest download first; stable repo_id tie-break for equal/missing mtimes.
cached = sorted(
seen_lower.values(),
key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()),
)
return {"cached": cached}
except Exception as e:
logger.error(f"Error listing cached GGUF repos: {e}", exc_info = True)
@ -2447,18 +2538,39 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject)
)
if not has_weights:
continue
last_modified = max(
(
_blob_mtime(f)
for rev in repo_info.revisions
for f in rev.files
if f.file_name.endswith(_WEIGHT_EXTENSIONS)
),
default = 0.0,
)
key = repo_id.lower()
existing = seen_lower.get(key)
if existing is None or total_size > existing["size_bytes"]:
seen_lower[key] = {
row = {
"repo_id": repo_id,
"size_bytes": total_size,
}
# Keep the newest timestamp across duplicate caches;
# attach only when known so absent rows sort as oldest.
lm = max(last_modified, (existing or {}).get("last_modified", 0.0))
if lm > 0:
row["last_modified"] = lm
seen_lower[key] = row
elif last_modified > existing.get("last_modified", 0.0):
existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
continue
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
# Newest download first; stable repo_id tie-break for equal/missing mtimes.
cached = sorted(
seen_lower.values(),
key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()),
)
return {"cached": cached}
except Exception as e:
logger.error(f"Error listing cached models: {e}", exc_info = True)

31
studio/backend/routes/providers.py
View file

@ -190,7 +190,8 @@ async def test_provider(
"""
Test connectivity to an external provider.
Makes a lightweight GET /models call to verify the API key works.
Makes a lightweight GET /models call to verify the API key works. Generic
custom endpoints use a chat-completions probe because /models is optional.
encrypted_api_key is decrypted server-side and never stored.
"""
info = get_provider_info(payload.provider_type)
@ -212,6 +213,14 @@ async def test_provider(
)
base_url = payload.base_url or info["base_url"]
if payload.provider_type == "custom":
if not base_url:
return ProviderTestResult(
success = False,
message = "Connection failed: Base URL is required for custom providers.",
models_count = None,
)
client = ExternalProviderClient(
provider_type = payload.provider_type,
base_url = base_url,
@ -220,6 +229,26 @@ async def test_provider(
)
try:
if payload.provider_type == "custom":
model_id = (payload.model_id or "").strip()
if not model_id:
return ProviderTestResult(
success = False,
message = "Connection failed: add a model ID to test custom providers.",
models_count = None,
)
await client.chat_completion(
messages = [{"role": "user", "content": "ping"}],
model = model_id,
temperature = 0.0,
top_p = 1.0,
max_tokens = 1,
)
return ProviderTestResult(
success = True,
message = "Connected successfully. Chat completions endpoint responded.",
models_count = None,
)
if info.get("model_list_mode") == "curated":
await client.verify_models_endpoint_lightweight()
return ProviderTestResult(

64
studio/backend/routes/rag.py
View file

@ -75,6 +75,19 @@ def _save_upload(file: UploadFile) -> tuple[str, str]:
return stored_path, filename
def _remove_stored_upload(stored_path: str | None) -> None:
"""Best-effort cleanup for files saved by _save_upload."""
if not stored_path:
return
try:
uploads = os.path.realpath(str(rag_uploads_root()))
target = os.path.realpath(stored_path)
if os.path.isfile(target) and os.path.commonpath([uploads, target]) == uploads:
os.remove(target)
except Exception: # noqa: BLE001 - DB/index deletion has already succeeded.
logger.warning("failed to remove RAG upload %s", stored_path, exc_info = True)
def _doc_view(row: dict) -> dict:
return {
"id": row["id"],
@ -84,6 +97,7 @@ def _doc_view(row: dict) -> dict:
"numChunks": row.get("num_chunks") or 0,
"kbId": row.get("kb_id"),
"threadId": row.get("thread_id"),
"projectId": row.get("project_id"),
"createdAt": row.get("created_at"),
}
@ -102,6 +116,7 @@ class SearchRequest(BaseModel):
query: str
kb_id: str | None = None
thread_id: str | None = None
project_id: str | None = None
top_k: int = Field(default = config.TOP_K_HYBRID, ge = 1, le = 50)
min_score: float = 0.0
mode: str = "hybrid" # hybrid | lexical | dense
@ -244,14 +259,50 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub
conn.close()
@router.post("/projects/{project_id}/documents")
async def upload_project_document(
project_id: str,
file: UploadFile = File(...),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
from storage.studio_db import get_chat_project
if get_chat_project(project_id) is None:
raise HTTPException(status_code = 404, detail = "Project not found")
stored_path, filename = _save_upload(file)
document_id, job_id = ingestion.start_ingestion(
store.project_scope(project_id),
None,
None,
filename,
stored_path,
project_id = project_id,
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@router.get("/projects/{project_id}/documents")
def list_project_documents(project_id: str, subject: str = Depends(get_current_subject)) -> dict:
_require_rag()
conn = rag_db.get_connection()
try:
docs = store.list_documents(conn, store.project_scope(project_id))
return {"documents": [_doc_view(d) for d in docs]}
finally:
conn.close()
@router.delete("/documents/{document_id}")
def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict:
_require_rag()
conn = rag_db.get_connection()
try:
if store.get_document(conn, document_id) is None:
doc = store.get_document(conn, document_id)
if doc is None:
raise HTTPException(status_code = 404, detail = "Document not found")
store.delete_document(conn, document_id)
_remove_stored_upload(doc.get("stored_path"))
return {"ok": True}
finally:
conn.close()
@ -297,10 +348,15 @@ def search(payload: SearchRequest, subject: str = Depends(get_current_subject))
_require_rag()
if payload.kb_id:
scope = store.kb_scope(payload.kb_id)
elif payload.thread_id:
scope = store.thread_scope(payload.thread_id)
else:
raise HTTPException(status_code = 400, detail = "Provide kb_id or thread_id")
scopes = []
if payload.project_id:
scopes.append(store.project_scope(payload.project_id))
if payload.thread_id:
scopes.append(store.thread_scope(payload.thread_id))
if not scopes:
raise HTTPException(status_code = 400, detail = "Provide kb_id, project_id, or thread_id")
scope = scopes[0] if len(scopes) == 1 else scopes
conn = rag_db.get_connection()
try:

19
studio/backend/routes/training.py
View file

@ -125,6 +125,17 @@ async def start_training(
backend = get_training_backend()
# S3 dataset loading needs the optional boto3 dependency. Reject early
# with a clear message so credentials are never accepted and then
# silently dropped on a host without boto3 installed.
if request.s3_config is not None:
from core.training.s3_dataset import boto3_available
if not boto3_available():
raise HTTPException(
status_code = 501,
detail = "S3 dataset loading requires boto3. Install it with: pip install boto3",
)
# Check before mutating state.
if backend.is_training_active():
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
@ -204,6 +215,9 @@ async def start_training(
"save_steps": request.save_steps,
"weight_decay": request.weight_decay,
"max_grad_norm": request.max_grad_norm,
"max_grad_value": request.max_grad_value,
"max_grad_leaf_norm": request.max_grad_leaf_norm,
"cast_norm_output_to_input_dtype": request.cast_norm_output_to_input_dtype,
"random_seed": request.random_seed,
"packing": request.packing,
"optim": request.optim,
@ -235,6 +249,7 @@ async def start_training(
"resume_from_checkpoint": request.resume_from_checkpoint,
"trust_remote_code": request.trust_remote_code,
"gpu_ids": request.gpu_ids,
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
}
# Training page has no trust_remote_code toggle; as a safety net consult
@ -293,6 +308,10 @@ async def start_training(
error = None,
)
except HTTPException:
# Deliberate rejections (S3 not implemented, resume validation) must
# reach the client with their original status, not a generic 500.
raise
except ValueError as e:
logger.warning("Rejected training GPU selection: %s", e)
# Deliberate user-facing GPU-selection validation message.

18
studio/backend/run.py
View file

@ -545,6 +545,11 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth`
# does) so its lazy submodule imports (export, hardware, mlx) and the
# DiffusionGemma runner never trip the install guard on a clean install.
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
def _write_pid_file():
"""Write the current process PID to the studio PID file."""
@ -864,6 +869,13 @@ def run_server(
from main import app, setup_frontend, _IS_COLAB
from utils.paths import ensure_studio_directories
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
# but never on Colab, which is a hosted VM reachable through its proxy. The
# gate reads the env var at request time, so this need not precede the import.
from utils.host_policy import apply_stdio_mcp_loopback_default
apply_stdio_mcp_loopback_default(host, is_colab = _IS_COLAB)
# Create all standard directories on startup.
ensure_studio_directories()
@ -1027,9 +1039,13 @@ def run_server(
_cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel
from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
_cloudflare_url = start_studio_tunnel(port)
app.state.cloudflare_url = _cloudflare_url
# Backstop: tear the tunnel down even on an abnormal exit that bypasses
# _graceful_shutdown (e.g. an exception after startup -> sys.exit). Idempotent.
atexit.register(stop_studio_tunnel)
except Exception as e:
logger.debug("Cloudflare tunnel skipped: %s", e)

2
studio/backend/startup_banner.py
View file

@ -94,6 +94,8 @@ def print_studio_access_banner(
external_url = f"http://{display_host}:{port}"
listen_all = bind_host in ("0.0.0.0", "::")
# The exact aliases the canned loopback_url below is valid for; any other bind
# (e.g. a specific LAN IP) must show its real address, not http://127.0.0.1.
loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1")
# Use the loopback URL only when reachable on loopback; otherwise show

139
studio/backend/state/tool_approvals.py Normal file
View file

@ -0,0 +1,139 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Per-call tool-call confirmation gate.
When a chat request sets ``confirm_tool_calls``, the agentic loop pauses
before executing each tool and waits here for the user's decision, which
arrives via ``POST /api/inference/tool-confirm`` on a separate connection.
Each gated call is identified by a unique ``approval_id`` (minted with
``new_approval_id``) that the loop both registers here and echoes in the
``tool_start`` stream event. The frontend sends that exact id back, so a
stale or duplicate confirmation -- or a second tool awaiting a decision in
the same session -- can never resolve the wrong call. ``session_id`` is
kept alongside purely as a scope check.
The slot is registered with ``begin_tool_decision`` *before* the loop
yields ``tool_start``, closing the race where a fast confirmation (or an
auto "Always allow") could otherwise arrive before the waiter exists.
``wait_tool_decision`` then blocks and cleans up its own slot.
"""
import secrets
import threading
from typing import Optional
# Generous ceiling so a user can deliberate; cancellation (stop button /
# disconnect) still breaks the wait early via ``cancel_event``.
_DECISION_TIMEOUT = 3600.0
# Fed to the model as the tool result when the user denies a call, so it
# can adapt and keep responding instead of the turn ending abruptly.
TOOL_REJECTED_MESSAGE = "The user declined to run this tool call."
_lock = threading.Lock()
# approval_id -> {"event": threading.Event, "decision": str|None, "session": str}
_pending: dict[str, dict] = {}
def new_approval_id() -> str:
"""Mint an unguessable id for one pending tool-call confirmation."""
return secrets.token_urlsafe(16)
def begin_tool_decision(session_id, approval_id) -> dict:
"""Register a pending decision slot and return it.
Call this *before* yielding the ``tool_start`` event so the waiter
always exists by the time the user's confirmation can arrive.
"""
slot = {
"event": threading.Event(),
"decision": None,
"session": session_id or "",
}
with _lock:
_pending[approval_id] = slot
return slot
def wait_tool_decision(
slot,
approval_id,
cancel_event = None,
timeout = _DECISION_TIMEOUT,
):
"""Block on a slot from ``begin_tool_decision`` until the user decides.
Returns ``"allow"`` or ``"deny"``. Falls back to ``"deny"`` if the wait
times out or generation is cancelled before the user decides. Always
removes its own slot on exit.
"""
try:
waited = 0.0
while not slot["event"].wait(timeout = 0.5):
if cancel_event is not None and cancel_event.is_set():
return "deny"
waited += 0.5
if waited >= timeout:
return "deny"
return slot["decision"] or "deny"
finally:
with _lock:
if _pending.get(approval_id) is slot:
_pending.pop(approval_id, None)
def abort_tool_decision(slot, approval_id) -> None:
"""Remove a slot that was announced but never entered ``wait_tool_decision``.
Streaming wrappers may stop after ``tool_start`` is yielded and before
the loop resumes into ``wait_tool_decision``. In that case there is no
waiter to run the normal cleanup path, so the generator close path calls
this explicitly.
"""
with _lock:
if _pending.get(approval_id) is slot:
_pending.pop(approval_id, None)
def request_tool_decision(
session_id,
approval_id,
cancel_event = None,
timeout = _DECISION_TIMEOUT,
):
"""Register and wait in one call (when the slot is not needed early)."""
slot = begin_tool_decision(session_id, approval_id)
return wait_tool_decision(slot, approval_id, cancel_event = cancel_event, timeout = timeout)
def resolve_tool_decision(
approval_id,
decision,
session_id = None,
) -> bool:
"""Record the user's "allow"/"deny" decision and unblock the loop.
Returns ``True`` if a pending call matched, ``False`` otherwise (e.g. a
stale or duplicate confirmation, or a session-scope mismatch).
The first decision wins: once a slot's event is set, a later (duplicate or
out-of-order) confirmation for the same id is rejected without mutating the
recorded decision, so an Allow can never be flipped to Deny in the window
before the waiter reads ``slot["decision"]`` and pops the slot.
"""
if not approval_id:
return False
with _lock:
slot = _pending.get(approval_id)
if not slot:
return False
if session_id is not None and slot["session"] != (session_id or ""):
return False
if slot["event"].is_set():
return False
slot["decision"] = decision
slot["event"].set()
return True

5
studio/backend/storage/rag_db.py
View file

@ -57,6 +57,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
scope TEXT NOT NULL,
kb_id TEXT,
thread_id TEXT,
project_id TEXT,
filename TEXT NOT NULL,
sha256 TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'pending',
@ -102,6 +103,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
);
"""
)
# Lazy upgrade for databases created before project sources existed.
cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()}
if "project_id" not in cols:
conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT")
def get_connection() -> sqlite3.Connection:

2
studio/backend/storage/studio_db.py
View file

@ -652,7 +652,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
r.ended_at, r.total_steps, r.final_step, r.final_loss,
r.output_dir, r.duration_seconds, r.error_message,
r.loss_sparkline, r.display_name,
r.loss_sparkline, r.display_name, r.config_json,
CASE
WHEN r.status = 'stopped'
AND r.output_dir IS NOT NULL

13
studio/backend/tests/test_anthropic_messages.py
View file

@ -1543,6 +1543,19 @@ class TestAnthropicMessagesToolRouting:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
def test_confirm_tool_calls_rejected_for_server_tools(self, monkeypatch):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(
confirm_tool_calls = True,
tools = [{"type": "web_search_20250305", "name": "web_search"}],
)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"]
assert backend.calls == []
def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(

74
studio/backend/tests/test_apple_gpu_sensors.py Normal file
View file

@ -0,0 +1,74 @@
# 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 Apple Silicon GPU sensors (SMC temperature + IOReport power)."""
import ctypes
import platform
import time
import pytest
from utils.hardware import apple
_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
class TestFourcc:
def test_roundtrip(self):
for key in ("#KEY", "Tg0D", "flt "):
assert apple._fourcc_str(apple._fourcc(key)) == key
def test_known_value(self):
# "flt " FourCC, same constant macmon uses.
assert apple._fourcc("flt ") == 1718383648
class TestWatts:
def test_millijoules(self):
assert apple._watts(2000, "mJ", 2.0) == pytest.approx(1.0)
def test_microjoules(self):
assert apple._watts(5_000_000, "uJ", 1.0) == pytest.approx(5.0)
def test_nanojoules(self):
assert apple._watts(1_500_000_000, "nJ", 1.0) == pytest.approx(1.5)
def test_unknown_unit_returns_none(self):
assert apple._watts(1000, "J", 1.0) is None
def test_zero_elapsed_returns_none(self):
assert apple._watts(1000, "mJ", 0.0) is None
class TestAverageValidTemps:
def test_averages_and_rounds(self):
assert apple._average_valid_temps([40.0, 50.0, 60.05]) == 50.0
def test_filters_invalid(self):
assert apple._average_valid_temps([-1.0, 0.0, 151.0, 42.0]) == 42.0
def test_empty_returns_none(self):
assert apple._average_valid_temps([]) is None
assert apple._average_valid_temps([0.0, 200.0]) is None
class TestSmcStructLayout:
def test_key_data_matches_smc_protocol_size(self):
# The AppleSMC user client rejects calls whose struct size differs.
assert ctypes.sizeof(apple._SMCKeyData) == 80
@pytest.mark.skipif(not _IS_APPLE_SILICON, reason = "requires Apple Silicon")
class TestLiveSensors:
def test_gpu_temperature_in_plausible_range(self):
temp = apple.read_gpu_temperature_c()
assert temp is not None
assert 0.0 < temp <= 150.0
def test_gpu_power_after_baseline(self):
apple.read_gpu_power_w() # first call only sets the baseline
time.sleep(0.3)
power = apple.read_gpu_power_w()
assert power is not None
assert power >= 0.0

153
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -371,3 +371,156 @@ def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeyp
"cache_path": str(vision_repo.repo_path),
}
]
def _gfile(name: str, size: int, mtime: float) -> SimpleNamespace:
"""A cached file carrying a Hugging Face ``blob_last_modified`` timestamp."""
return SimpleNamespace(
file_name = name,
size_on_disk = size,
blob_path = None,
blob_last_modified = mtime,
)
def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_path):
"""An unreadable auxiliary cache (e.g. an inaccessible
``~/.cache/huggingface/hub``) must be skipped, not abort the scan.
Regression guard for ``extra.is_dir()`` raising and wiping the response.
"""
import huggingface_hub
import utils.paths as paths_mod
active = SimpleNamespace(
repos = [_repo("Org/Active", [_file("Q4_K_M.gguf", 5_000)], tmp_path / "active")]
)
def _fake_scan(cache_dir = None):
if cache_dir is None:
return active
raise AssertionError("auxiliary scan should have been skipped")
class _Boom:
def is_dir(self):
raise PermissionError(13, "Permission denied")
def resolve(self):
raise PermissionError(13, "Permission denied")
monkeypatch.setattr(huggingface_hub, "scan_cache_dir", _fake_scan)
monkeypatch.setattr(paths_mod, "legacy_hf_cache_dir", lambda: _Boom())
monkeypatch.setattr(paths_mod, "hf_default_cache_dir", lambda: _Boom())
scans = models_route._all_hf_cache_scans()
assert scans == [active]
# End-to-end: the endpoint still returns the active cache's repo.
monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [active])
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/Active",
"size_bytes": 5_000,
"cache_path": str(tmp_path / "active"),
}
]
def test_list_cached_gguf_sorts_newest_first_grouping_by_latest_quant(monkeypatch, tmp_path):
"""Downloaded is ordered newest-first, and a multi-quant repo is placed by
its most recently downloaded quant (``last_modified`` = newest quant)."""
older = _repo(
"Org/Older",
[_gfile("Older-Q4_K_M.gguf", 5_000, 1_000.0)],
tmp_path / "models--Org--Older",
)
newer = _repo(
"Org/Newer",
[
_gfile("Newer-Q4_K_M.gguf", 5_000, 2_000.0),
_gfile("Newer-Q8_0.gguf", 9_000, 3_000.0), # newest quant in the repo
],
tmp_path / "models--Org--Newer",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [older, newer])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert [c["repo_id"] for c in result["cached"]] == ["Org/Newer", "Org/Older"]
assert result["cached"][0]["last_modified"] == 3_000.0
assert result["cached"][1]["last_modified"] == 1_000.0
def test_list_cached_gguf_dedupe_keeps_newest_timestamp(monkeypatch, tmp_path):
"""Same repo in two caches with equal size keeps the newest last_modified,
regardless of scan order."""
older = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 1_000.0)], tmp_path / "a")
newer = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 9_000.0)], tmp_path / "b")
for scans in ([older, newer], [newer, older]): # both orders
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda s = scans: [SimpleNamespace(repos = [s[0]]), SimpleNamespace(repos = [s[1]])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "t"))
assert len(result["cached"]) == 1
assert result["cached"][0]["last_modified"] == 9_000.0
def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_path):
"""The per-quant 'downloaded' flag is driven by the real weight file in a
single snapshot; an mmproj vision adapter (matching a quant label) must
not make that quant appear downloaded."""
import huggingface_hub.constants as hf_constants
variants = [
SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000),
SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000),
]
monkeypatch.setattr(
models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True)
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16"
result = asyncio.run(
models_route.get_gguf_variants(
repo_id = "org/repo", hf_token = None, current_subject = "test-user"
)
)
flags = {v.quant: v.downloaded for v in result.variants}
assert flags["Q4_K_M"] is True
assert flags["F16"] is False
def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path):
"""A cached mmproj adapter must not count toward a same-label main
variant's download progress (mmproj-F16 vs an F16 weight)."""
import huggingface_hub.constants as hf_constants
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # only the adapter on disk
result = asyncio.run(
models_route.get_gguf_download_progress(
repo_id = "org/repo",
variant = "F16",
expected_bytes = 20_000,
current_subject = "test-user",
)
)
assert result["downloaded_bytes"] == 0
assert result["progress"] == 0

304
studio/backend/tests/test_cloudflare_tunnel.py
View file

@ -52,6 +52,26 @@ def test_url_regex_no_match_on_unrelated():
assert ct._URL_RE.search("INF connecting to https://api.cloudflare.com/v4") is None
def test_url_regex_ignores_api_endpoint():
# cloudflared's failure line names its own API host; it must never be taken
# as the tunnel URL (it returns a 404 and is not a quick tunnel).
line = (
'failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel": '
"context deadline exceeded"
)
assert ct._URL_RE.search(line) is None
def test_url_regex_skips_api_host_but_matches_real_url():
blob = (
'ERR failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel"\n'
"INF | https://brave-mountain-river-clouds.trycloudflare.com |\n"
)
m = ct._URL_RE.search(blob)
assert m is not None
assert m.group(0) == "https://brave-mountain-river-clouds.trycloudflare.com"
# ── asset mapping ────────────────────────────────────────────────────
@ -295,9 +315,88 @@ def test_stop_terminates_process():
t.stop()
def test_wait_for_url_times_out_without_blocking():
def test_start_after_stop_does_not_spawn(monkeypatch):
# If stop() lands before start() (a concurrent shutdown in the caller's
# register->start window), start() must NOT spawn a cloudflared process --
# nobody would own it and it would be orphaned.
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
assert t.wait_for_url(timeout = 0.05) is None
spawned = []
class _FakeProc:
stdout = None
def poll(self):
return 0
monkeypatch.setattr(ct.subprocess, "Popen", lambda *a, **k: (spawned.append(a), _FakeProc())[1])
t.stop() # proc is None -> no-op terminate, but marks the tunnel stopped
t.start() # must short-circuit before Popen
assert spawned == []
assert t._proc is None
def test_wait_for_ready_times_out_without_blocking():
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
assert t.wait_for_ready(timeout = 0.05) is None
def _fake_proc(text):
return types.SimpleNamespace(stdout = io.StringIO(text))
def test_reader_captures_url_and_registration():
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
t._reader(
_fake_proc(
"INF Requesting new quick Tunnel on trycloudflare.com...\n"
"INF | https://words-here-abc.trycloudflare.com |\n"
"INF Registered tunnel connection connIndex=0 protocol=http2\n"
)
)
assert t.url == "https://words-here-abc.trycloudflare.com"
assert t.ready is True
assert t.wait_for_ready(0) == t.url
assert t.error is None # a fully-registered tunnel records no error
def test_reader_url_without_registration_is_not_ready():
# A URL but no "Registered tunnel connection" (e.g. quic control stream
# fails) must not be advertised -- it returns Cloudflare error 1033.
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
t._reader(
_fake_proc(
"INF | https://words-here-abc.trycloudflare.com |\n"
'ERR failed to serve tunnel connection error="control stream failure"\n'
)
)
assert t.url == "https://words-here-abc.trycloudflare.com"
assert t.ready is False
assert t.wait_for_ready(0) is None
assert t.error == "cloudflared exited before the tunnel connection registered"
def test_reader_handles_none_stdout():
# Popen.stdout can be None; _reader must not crash and must leave the tunnel
# un-ready so wait_for_ready returns None.
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
t._reader(types.SimpleNamespace(stdout = None))
assert t.url is None
assert t.ready is False
assert t.wait_for_ready(0) is None
assert t.error == "cloudflared exited before emitting a tunnel URL"
def test_reader_ignores_api_endpoint_failure_line():
t = ct.CloudflareTunnel(8080, "/bin/cloudflared")
t._reader(
_fake_proc(
"ERR failed to request quick Tunnel: Post "
'"https://api.trycloudflare.com/tunnel": context deadline exceeded\n'
)
)
assert t.url is None
assert t.wait_for_ready(0) is None
assert t.error == "cloudflared exited before emitting a tunnel URL"
def test_start_studio_tunnel_no_binary(monkeypatch):
@ -306,18 +405,23 @@ def test_start_studio_tunnel_no_binary(monkeypatch):
def test_start_studio_tunnel_registers_before_wait(monkeypatch):
# The tunnel must be visible to stop_studio_tunnel() during the URL wait,
# else a shutdown in that window orphans cloudflared.
# The tunnel must be visible to stop_studio_tunnel() during the readiness
# wait, else a shutdown in that window orphans cloudflared.
seen = {}
class _Stub:
def __init__(self, port, binary):
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
def start(self):
pass
def wait_for_url(self, timeout):
def wait_for_ready(self, timeout):
seen["active_during_wait"] = ct._active_tunnel is self
self.url = "https://x.trycloudflare.com"
return self.url
@ -338,13 +442,18 @@ def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch):
seen = {}
class _Stub:
def __init__(self, port, binary):
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
def start(self):
pass
def wait_for_url(self, timeout):
def wait_for_ready(self, timeout):
return None
def stop(self):
@ -359,13 +468,18 @@ def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch):
def test_start_studio_tunnel_returns_url(monkeypatch):
class _StubTunnel:
def __init__(self, port, binary):
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
def start(self):
self.url = "https://stub-xyz.trycloudflare.com"
def wait_for_url(self, timeout):
def wait_for_ready(self, timeout):
return self.url
def stop(self):
@ -379,6 +493,169 @@ def test_start_studio_tunnel_returns_url(monkeypatch):
ct.stop_studio_tunnel()
def test_start_studio_tunnel_falls_back_to_http2(monkeypatch):
# First attempt mints a URL but never registers (quic blocked); the http2
# retry registers and wins.
attempts = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.protocol = protocol
self.url = None
attempts.append(protocol)
def start(self):
self.url = "https://words.trycloudflare.com" # URL always minted
def wait_for_ready(self, timeout):
return self.url if self.protocol == "http2" else None
def stop(self):
pass
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
try:
assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com"
assert attempts == [None, "http2"] # default first, then forced http2
finally:
ct.stop_studio_tunnel()
def test_start_studio_tunnel_no_retry_when_shutdown_between_attempts(monkeypatch):
# A stop() landing in the gap AFTER the failed first attempt is cleaned up but
# BEFORE the http2 retry registers must abort the loop -- not start a second
# tunnel that nobody will ever stop (Codex review). Simulated by having the
# first attempt's stop() (called during cleanup) trigger the shutdown.
attempts = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
attempts.append(protocol)
def start(self):
self.url = "https://words.trycloudflare.com" # URL minted, never ready
def wait_for_ready(self, timeout):
return None
def stop(self):
ct.stop_studio_tunnel() # a concurrent shutdown lands in the gap
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
assert ct.start_studio_tunnel(8080) is None
assert attempts == [None] # http2 retry aborted after shutdown
assert ct._active_tunnel is None
def test_start_studio_tunnel_no_http2_retry_when_no_url(monkeypatch):
# No URL at all is an API/network failure; the http2 fallback would not help,
# so it must be skipped (don't burn a second timeout window).
attempts = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
attempts.append(protocol)
def start(self):
pass # never mints a URL
def wait_for_ready(self, timeout):
return None
def stop(self):
pass
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
assert ct.start_studio_tunnel(8080) is None
assert attempts == [None]
def test_start_studio_tunnel_both_protocols_fail_registration(monkeypatch):
# Both quic and http2 mint a URL but neither registers -> both attempts are
# exhausted and None is returned (no dead URL advertised).
attempts = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
attempts.append(protocol)
def start(self):
self.url = "https://words.trycloudflare.com" # URL minted, never ready
def wait_for_ready(self, timeout):
return None
def stop(self):
pass
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
assert ct.start_studio_tunnel(8080) is None
assert attempts == [None, "http2"]
assert ct._active_tunnel is None
def test_start_studio_tunnel_aborts_retry_on_concurrent_shutdown(monkeypatch):
# If a concurrent stop_studio_tunnel() clears _active_tunnel while we wait,
# the retry loop must NOT start a second (http2) tunnel: shutdown is already
# done, so nothing would ever stop it and it would be orphaned.
attempts = []
class _Stub:
def __init__(
self,
port,
binary,
protocol = None,
):
self.url = None
attempts.append(protocol)
def start(self):
self.url = "https://words.trycloudflare.com" # URL minted (saw_url True)
def wait_for_ready(self, timeout):
# Simulate stop_studio_tunnel() landing during the wait.
with ct._active_lock:
ct._active_tunnel = None
return None # never registered
def stop(self):
pass
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
assert ct.start_studio_tunnel(8080) is None
assert attempts == [None] # no http2 retry -> no orphaned second tunnel
assert ct._active_tunnel is None
# ── run.py source-level pins (AST / source, no heavy import) ─────────
@ -419,6 +696,13 @@ def test_argparse_cloudflare_default_true():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True
def test_run_server_registers_tunnel_atexit_backstop():
# An abnormal exit (exception after startup -> sys.exit) bypasses
# _graceful_shutdown; an atexit backstop must still stop the tunnel.
src = _RUN_PY.read_text()
assert "atexit.register(stop_studio_tunnel)" in src
def test_run_server_gates_tunnel_on_wildcard():
# Guard against accidentally widening the trigger beyond 0.0.0.0.
source = _RUN_PY.read_text()

278
studio/backend/tests/test_datacenter_gpu_tuning.py Normal file
View file

@ -0,0 +1,278 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Data-center llama.cpp env tuning: FP32 accum (+ P2P / launch queues for
multi-GPU) must apply only to datacenter NVIDIA parts, never consumer GeForce,
AMD/ROCm, CPU or macOS. User values win; UNSLOTH_DISABLE_DC_TUNING=1 disables.
"""
from __future__ import annotations
import sys
import types
import pytest
from core.inference.llama_cpp import LlamaCppBackend
def _fake_torch(
names,
*,
hip = None,
cuda_ok = True,
):
"""torch stub: version.hip, cuda.is_available/device_count, get_device_properties(i).name."""
t = types.ModuleType("torch")
t.version = types.SimpleNamespace(hip = hip)
t.cuda = types.SimpleNamespace(
is_available = lambda: cuda_ok,
device_count = lambda: len(names),
get_device_properties = lambda i: types.SimpleNamespace(name = names[i]),
)
return t
@pytest.fixture(autouse = True)
def _clear_cuda_visible_devices(monkeypatch):
"""Detection reads CUDA_VISIBLE_DEVICES, so clear it by default (run unmasked,
physical id == ordinal) regardless of host; masked tests set it explicitly."""
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
# ---------------------------------------------------------------------------
# _is_datacenter_gpu
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"names,expected",
[
# Datacenter / professional parts.
(["NVIDIA A100-SXM4-80GB"], True),
(["NVIDIA A30"], True),
(["NVIDIA H100 80GB HBM3"], True),
(["NVIDIA H200"], True),
(["NVIDIA H800"], True),
(["NVIDIA GH200 480GB"], True),
(["NVIDIA B200"], True),
(["NVIDIA GB200"], True),
(["NVIDIA L40S"], True),
(["NVIDIA L4"], True),
(["NVIDIA RTX PRO 6000 Blackwell Server Edition"], True),
(["NVIDIA RTX 6000 Ada Generation"], True),
# Consumer GeForce: never.
(["NVIDIA GeForce RTX 4090"], False),
(["NVIDIA GeForce RTX 5090"], False),
(["NVIDIA GeForce RTX 3090"], False),
(["NVIDIA GeForce RTX 2080 Ti"], False),
(["NVIDIA GeForce GTX 1080"], False),
# Workstation/laptop: short markers must not match as substrings
# ("a100" in "A1000", "a30" in "A3000").
(["NVIDIA RTX A1000 Laptop GPU"], False),
(["NVIDIA RTX A1000 6GB Laptop GPU"], False),
(["NVIDIA RTX A3000 Laptop GPU"], False),
# Homogeneous multi-DC: all must match.
(["NVIDIA B200", "NVIDIA B200"], True),
(["NVIDIA H100 80GB HBM3", "NVIDIA H100 80GB HBM3"], True),
# Mixed DC + consumer: non-DC, so tuning never lands on the GeForce.
(["NVIDIA B200", "NVIDIA GeForce RTX 4090"], False),
(["NVIDIA GeForce RTX 4090", "NVIDIA B200"], False),
],
)
def test_is_datacenter_gpu(monkeypatch, names, expected):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(names))
assert LlamaCppBackend._is_datacenter_gpu() is expected
def test_is_datacenter_gpu_respects_selection(monkeypatch):
# A mixed box where only the DC GPU is selected -> True; only consumer -> False.
monkeypatch.setitem(
sys.modules,
"torch",
_fake_torch(["NVIDIA B200", "NVIDIA GeForce RTX 4090"]),
)
assert LlamaCppBackend._is_datacenter_gpu([0]) is True
assert LlamaCppBackend._is_datacenter_gpu([1]) is False
assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False
def test_is_datacenter_gpu_out_of_range_indices_skipped(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"]))
# Out-of-range / negative indices are skipped; the one valid DC GPU still wins.
assert LlamaCppBackend._is_datacenter_gpu([0, 5, -1]) is True
# Only invalid indices -> nothing seen -> False (fail closed for the flag).
assert LlamaCppBackend._is_datacenter_gpu([5, 9]) is False
def test_is_datacenter_gpu_masked_host_physical_ids(monkeypatch):
# Mask 4,5,6,7 -> ordinals 0..3 == physical 4..7. PHYSICAL selection [4,5]
# must resolve, not index out of range (the pre-fix bug: 4 >= device_count).
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7")
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4))
assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is True
assert LlamaCppBackend._is_datacenter_gpu([4, 5, 6, 7]) is True
assert LlamaCppBackend._is_datacenter_gpu(None) is True
assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False # not visible -> skip
def test_is_datacenter_gpu_masked_host_reordered(monkeypatch):
# Reordered mask preserves order: ordinal 0 -> physical 7, 1 -> 4, ...
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,4,5,6")
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100 80GB HBM3"] * 4))
assert LlamaCppBackend._is_datacenter_gpu([7, 4]) is True
def test_is_datacenter_gpu_masked_host_mixed_class(monkeypatch):
# Mask 4,5: physical 4 = GeForce, physical 5 = B200. Detection must follow the
# selected physical GPU, not a same-numbered ordinal.
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5")
monkeypatch.setitem(
sys.modules,
"torch",
_fake_torch(["NVIDIA GeForce RTX 4090", "NVIDIA B200"]),
)
assert LlamaCppBackend._is_datacenter_gpu([4]) is False
assert LlamaCppBackend._is_datacenter_gpu([5]) is True
assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is False
def test_is_datacenter_gpu_unparsable_mask_falls_back(monkeypatch):
# Unparsable (UUID) mask falls back to physical id == ordinal (mirrors
# _get_gpu_free_memory), so ordinal lookup still classifies the device.
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-abcdef12")
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"]))
assert LlamaCppBackend._is_datacenter_gpu([0]) is True
def test_is_datacenter_gpu_rocm_is_false(monkeypatch):
# ROCm reuses torch.cuda.*; an MI300X must not qualify.
monkeypatch.setitem(
sys.modules,
"torch",
_fake_torch(["AMD Instinct MI300X"], hip = "6.2.0"),
)
assert LlamaCppBackend._is_datacenter_gpu() is False
def test_is_datacenter_gpu_no_cuda_is_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False))
assert LlamaCppBackend._is_datacenter_gpu() is False
def test_is_datacenter_gpu_missing_torch_is_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", None)
assert LlamaCppBackend._is_datacenter_gpu() is False
# ---------------------------------------------------------------------------
# _effective_gpu_count
# ---------------------------------------------------------------------------
def test_effective_gpu_count_explicit_selection(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4))
assert LlamaCppBackend._effective_gpu_count([0]) == 1
assert LlamaCppBackend._effective_gpu_count([0, 1, 2]) == 3
def test_effective_gpu_count_none_uses_visible(monkeypatch):
# None -> visible device count.
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4))
assert LlamaCppBackend._effective_gpu_count(None) == 4
def test_effective_gpu_count_no_cuda_is_zero(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False))
assert LlamaCppBackend._effective_gpu_count(None) == 0
def test_effective_gpu_count_missing_torch_is_zero(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", None)
assert LlamaCppBackend._effective_gpu_count(None) == 0
# ---------------------------------------------------------------------------
# _apply_datacenter_env (the env-injection decision)
# ---------------------------------------------------------------------------
def test_apply_env_single_dc_gpu_sets_only_fp32(monkeypatch):
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"]))
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, [0]) is True
assert env == {"GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "1"}
assert "GGML_CUDA_P2P" not in env # no multi-GPU flags on one GPU
assert "CUDA_SCALE_LAUNCH_QUEUES" not in env
def test_apply_env_multi_dc_gpu_sets_all(monkeypatch):
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4))
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True
assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1"
assert env["GGML_CUDA_P2P"] == "1"
assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x"
def test_apply_env_none_indices_uses_visible_count(monkeypatch):
# None on a 2x DC box -> multi-GPU flags applied.
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100", "NVIDIA H100"]))
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, None) is True
assert env["GGML_CUDA_P2P"] == "1"
assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x"
def test_apply_env_consumer_gpu_is_noop(monkeypatch):
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA GeForce RTX 4090"] * 2))
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False
assert env == {}
def test_apply_env_user_value_wins(monkeypatch):
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2))
env = {
"GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "0", # user explicitly disabled
"CUDA_SCALE_LAUNCH_QUEUES": "8x", # user override
}
assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True
# setdefault must not clobber user values; the unset one still defaults.
assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "0"
assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "8x"
assert env["GGML_CUDA_P2P"] == "1"
def test_apply_env_disable_flag_respected(monkeypatch):
monkeypatch.setenv("UNSLOTH_DISABLE_DC_TUNING", "1")
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2))
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False
assert env == {}
def test_apply_env_fail_open_on_detection_error(monkeypatch):
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setitem(sys.modules, "torch", None) # detection raises -> False
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, [0]) is False
assert env == {}
def test_apply_env_masked_host_multi_dc(monkeypatch):
# End-to-end masked host (mask 4,5,6,7, physical selection [4,5]): pre-fix
# applied no tuning; now all three multi-GPU flags must be set.
monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7")
monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4))
env: dict = {}
assert LlamaCppBackend._apply_datacenter_env(env, [4, 5]) is True
assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1"
assert env["GGML_CUDA_P2P"] == "1"
assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x"

65
studio/backend/tests/test_default_output_dir_name.py Normal file
View file

@ -0,0 +1,65 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Auto-generated training output dir names stay inside outputs_root.
Regression for local-model training: a model loaded by absolute path (e.g.
``G:\\modelsAI\\...\\gemma-4-12B-it`` on a non-system drive) used to seed the
default run dir with that full path, so ``resolve_output_dir`` raised
``path escapes root`` because the result was not under ``<studio>/outputs``.
"""
import importlib.util
from pathlib import Path
import pytest
_BACKEND_DIR = Path(__file__).resolve().parent.parent
def _load_storage_roots():
path = _BACKEND_DIR / "utils/paths/storage_roots.py"
spec = importlib.util.spec_from_file_location("storage_roots_under_test", path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def test_repo_id_keeps_namespace():
sr = _load_storage_roots()
assert sr.default_run_dir_name("unsloth/gemma-3-4b") == "unsloth_gemma-3-4b"
assert sr.default_run_dir_name("gemma-3-4b") == "gemma-3-4b"
def test_local_paths_collapse_to_basename():
sr = _load_storage_roots()
assert sr.default_run_dir_name(r"G:\modelsAI\gguf\test\gemma-4-12B-it") == "gemma-4-12B-it"
assert sr.default_run_dir_name("/data/models/gemma-3-4b") == "gemma-3-4b"
assert sr.default_run_dir_name("~/models/gemma-3-4b") == "gemma-3-4b"
assert sr.default_run_dir_name("C:/Users/me/models/gemma-3-4b") == "gemma-3-4b"
def test_empty_falls_back_to_model():
sr = _load_storage_roots()
assert sr.default_run_dir_name("") == "model"
assert sr.default_run_dir_name(" ") == "model"
def test_very_long_name_is_capped():
sr = _load_storage_roots()
name = sr.default_run_dir_name("a" * 500)
assert 0 < len(name) <= 200
def test_derived_name_resolves_under_outputs_root(tmp_path, monkeypatch):
sr = _load_storage_roots()
outputs = tmp_path / "outputs"
outputs.mkdir()
monkeypatch.setattr(sr, "outputs_root", lambda: outputs)
name = sr.default_run_dir_name(r"G:\modelsAI\gguf\test\gemma-4-12B-it")
resolved = sr.resolve_output_dir(f"{name}_1781327234")
assert resolved == outputs / "gemma-4-12B-it_1781327234"
# No escape: the absolute G: source no longer leaks into the output path.
assert "modelsAI" not in str(resolved)

472
studio/backend/tests/test_export_absolute_paths.py Normal file
View file

@ -0,0 +1,472 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import importlib.machinery
import importlib.util
import sys
import types
from pathlib import Path
import pytest
_BACKEND_DIR = Path(__file__).resolve().parent.parent
def _load_module(
module_name: str,
relative_path: str,
monkeypatch = None,
):
path = _BACKEND_DIR / relative_path
spec = importlib.util.spec_from_file_location(module_name, path)
assert spec is not None
assert spec.loader is not None
module = importlib.util.module_from_spec(spec)
if monkeypatch is None:
sys.modules[module_name] = module
else:
monkeypatch.setitem(sys.modules, module_name, module)
spec.loader.exec_module(module)
return module
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
class _Router:
def get(self, *args, **kwargs):
return lambda fn: fn
def post(self, *args, **kwargs):
return lambda fn: fn
def delete(self, *args, **kwargs):
return lambda fn: fn
class _HTTPException(Exception):
def __init__(
self,
status_code: int,
detail: str | None = None,
):
super().__init__(detail)
self.status_code = status_code
self.detail = detail
class _LocalModelInfo:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
def _identity_decorator(*_args, **_kwargs):
return lambda fn: fn
def _install_lightweight_backend_stubs(monkeypatch):
fastapi = types.ModuleType("fastapi")
fastapi.APIRouter = lambda: _Router()
fastapi.Body = lambda default = None, **_kwargs: default
fastapi.Depends = lambda dependency = None, **_kwargs: dependency
fastapi.HTTPException = _HTTPException
fastapi.Query = lambda default = None, **_kwargs: default
fastapi.Request = object
monkeypatch.setitem(sys.modules, "fastapi", fastapi)
fastapi_responses = types.ModuleType("fastapi.responses")
fastapi_responses.StreamingResponse = object
monkeypatch.setitem(sys.modules, "fastapi.responses", fastapi_responses)
monkeypatch.setitem(
sys.modules,
"structlog",
types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
),
)
loggers = types.ModuleType("loggers")
loggers.get_logger = lambda *args, **kwargs: _DummyLogger()
monkeypatch.setitem(sys.modules, "loggers", loggers)
auth_pkg = types.ModuleType("auth")
auth_mod = types.ModuleType("auth.authentication")
auth_mod.get_current_subject = lambda: None
monkeypatch.setitem(sys.modules, "auth", auth_pkg)
monkeypatch.setitem(sys.modules, "auth.authentication", auth_mod)
core_pkg = types.ModuleType("core")
core_export = types.ModuleType("core.export")
core_export.get_export_backend = lambda: None
core_inference = types.ModuleType("core.inference")
core_inference.get_inference_backend = lambda: None
monkeypatch.setitem(sys.modules, "core", core_pkg)
monkeypatch.setitem(sys.modules, "core.export", core_export)
monkeypatch.setitem(sys.modules, "core.inference", core_inference)
utils_pkg = types.ModuleType("utils")
utils_pkg.__path__ = []
utils_paths = types.ModuleType("utils.paths")
storage_roots = _load_module(
"utils.paths.storage_roots",
"utils/paths/storage_roots.py",
monkeypatch,
)
utils_pkg.paths = utils_paths
utils_paths.storage_roots = storage_roots
utils_paths.is_local_path = lambda value: Path(str(value)).is_absolute()
utils_paths.outputs_root = lambda: Path("outputs")
utils_paths.exports_root = storage_roots.exports_root
utils_paths.resolve_cached_repo_id_case = lambda value: value
utils_paths.resolve_output_dir = lambda value = None: Path(value or "outputs")
utils_paths.resolve_export_dir = storage_roots.resolve_export_dir
monkeypatch.setitem(sys.modules, "utils", utils_pkg)
monkeypatch.setitem(sys.modules, "utils.paths", utils_paths)
utils_utils = types.ModuleType("utils.utils")
utils_utils.log_and_http_error = lambda *args, **kwargs: (_ for _ in ()).throw(
_HTTPException(kwargs.get("status_code", 500), kwargs.get("detail"))
)
utils_utils.safe_error_detail = lambda value: str(value)
monkeypatch.setitem(sys.modules, "utils.utils", utils_utils)
utils_models = types.ModuleType("utils.models")
for name in (
"scan_trained_models",
"scan_exported_models",
"scan_checkpoints",
"list_gguf_variants",
):
setattr(utils_models, name, lambda *args, **kwargs: [])
for name in (
"get_base_model_from_checkpoint",
"get_base_model_from_lora",
"load_model_defaults",
):
setattr(utils_models, name, lambda *args, **kwargs: None)
utils_models.is_vision_model = lambda *args, **kwargs: False
utils_models.is_embedding_model = lambda *args, **kwargs: False
utils_models.ModelConfig = object
monkeypatch.setitem(sys.modules, "utils.models", utils_models)
utils_model_config = types.ModuleType("utils.models.model_config")
utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None
utils_model_config._extract_quant_label = lambda value: value
utils_model_config.is_audio_input_type = lambda *args, **kwargs: None
monkeypatch.setitem(
sys.modules,
"utils.models.model_config",
utils_model_config,
)
models_pkg = types.ModuleType("models")
models_pkg.__path__ = []
for name in (
"CheckpointInfo",
"CheckpointListResponse",
"LocalModelListResponse",
"ModelCheckpoints",
"ModelDetails",
"LoRAScanResponse",
"LoRAInfo",
"ModelListResponse",
"LoadCheckpointRequest",
"ExportStatusResponse",
"ExportOperationResponse",
"ExportMergedModelRequest",
"ExportBaseModelRequest",
"ExportGGUFRequest",
"ExportLoRAAdapterRequest",
):
setattr(models_pkg, name, object)
models_pkg.LocalModelInfo = _LocalModelInfo
monkeypatch.setitem(sys.modules, "models", models_pkg)
models_models = types.ModuleType("models.models")
for name in (
"BrowseEntry",
"BrowseFoldersResponse",
"GgufVariantDetail",
"GgufVariantsResponse",
"ScanFolderInfo",
"AddScanFolderRequest",
):
setattr(models_models, name, object)
models_models.ModelType = str
monkeypatch.setitem(sys.modules, "models.models", models_models)
models_responses = types.ModuleType("models.responses")
for name in (
"LoRABaseModelResponse",
"VisionCheckResponse",
"EmbeddingCheckResponse",
):
setattr(models_responses, name, object)
monkeypatch.setitem(sys.modules, "models.responses", models_responses)
def _install_pydantic_stub(monkeypatch):
pydantic = types.ModuleType("pydantic")
pydantic.BaseModel = object
pydantic.Field = lambda default = None, **_kwargs: default
pydantic.field_validator = _identity_decorator
monkeypatch.setitem(sys.modules, "pydantic", pydantic)
def _install_export_backend_stubs(monkeypatch):
_install_lightweight_backend_stubs(monkeypatch)
unsloth = types.ModuleType("unsloth")
unsloth.FastLanguageModel = object
unsloth.FastVisionModel = object
unsloth._IS_MLX = True
unsloth.__spec__ = importlib.machinery.ModuleSpec("unsloth", loader = None)
monkeypatch.setitem(sys.modules, "unsloth", unsloth)
unsloth_zoo = types.ModuleType("unsloth_zoo")
unsloth_zoo.__path__ = []
unsloth_zoo.__spec__ = importlib.machinery.ModuleSpec(
"unsloth_zoo",
loader = None,
is_package = True,
)
llama_cpp = types.ModuleType("unsloth_zoo.llama_cpp")
llama_cpp.LLAMA_CPP_DEFAULT_DIR = str(Path("/tmp/llama.cpp"))
llama_cpp._resolve_local_convert_script = lambda *args, **kwargs: None
llama_cpp.__spec__ = importlib.machinery.ModuleSpec(
"unsloth_zoo.llama_cpp",
loader = None,
)
monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo)
monkeypatch.setitem(sys.modules, "unsloth_zoo.llama_cpp", llama_cpp)
huggingface_hub = types.ModuleType("huggingface_hub")
huggingface_hub.HfApi = object
huggingface_hub.ModelCard = object
monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub)
utils_hardware = types.ModuleType("utils.hardware")
utils_hardware.clear_gpu_cache = lambda: None
monkeypatch.setitem(sys.modules, "utils.hardware", utils_hardware)
utils_models = sys.modules["utils.models"]
utils_models.get_base_model_from_lora = lambda *args, **kwargs: None
utils_models.is_vision_model = lambda *args, **kwargs: False
utils_model_config = sys.modules["utils.models.model_config"]
utils_model_config.detect_audio_type = lambda *args, **kwargs: None
utils_paths = sys.modules["utils.paths"]
utils_paths.ensure_dir = lambda path: Path(path).mkdir(parents = True, exist_ok = True)
utils_paths.resolve_export_write_dir = lambda value = None: Path(value or "exports")
utils_paths.resolve_output_dir = lambda value = None: Path(value or "outputs")
def test_gguf_export_cleans_temp_dir_when_post_processing_fails(tmp_path, monkeypatch):
_install_export_backend_stubs(monkeypatch)
export_mod = _load_module("test_core_export_backend", "core/export/export.py", monkeypatch)
cwd = tmp_path / "cwd"
save_dir = tmp_path / "export"
cwd.mkdir()
monkeypatch.chdir(cwd)
monkeypatch.setattr(export_mod, "resolve_export_write_dir", lambda _value: save_dir)
monkeypatch.setattr(
export_mod.shutil,
"move",
lambda *args, **kwargs: (_ for _ in ()).throw(OSError("move failed")),
)
class _Model:
def save_pretrained_gguf(self, model_save_path, tokenizer, quantization_method):
Path(model_save_path).mkdir(parents = True)
(Path(model_save_path) / "model.safetensors").write_bytes(b"weights")
(cwd / "converted.gguf").write_bytes(b"gguf")
backend = export_mod.ExportBackend.__new__(export_mod.ExportBackend)
backend.current_model = _Model()
backend.current_tokenizer = object()
backend.current_checkpoint = None
success, message, output_path = backend.export_gguf(str(save_dir), "Q4_K_M")
assert success is False
assert "move failed" in message
assert output_path is None
assert list(save_dir.glob("_tmp_model_*")) == []
def test_save_directory_validator_rejects_windows_parent_segments(monkeypatch):
_install_pydantic_stub(monkeypatch)
export_models = _load_module("test_models_export", "models/export.py", monkeypatch)
with pytest.raises(ValueError, match = r"\.\."):
export_models._validate_save_directory(r"E:\AI\..\secret")
def test_save_directory_validator_allows_deep_absolute_paths(monkeypatch, tmp_path):
_install_pydantic_stub(monkeypatch)
export_models = _load_module("test_models_export_deep_path", "models/export.py", monkeypatch)
deep_path = tmp_path
for index in range(40):
deep_path /= f"segment-{index:02d}"
raw = str(deep_path)
assert len(raw) > 255
assert export_models._validate_save_directory(raw) == raw
def test_save_directory_validator_rejects_long_path_component(monkeypatch, tmp_path):
_install_pydantic_stub(monkeypatch)
export_models = _load_module(
"test_models_export_long_component", "models/export.py", monkeypatch
)
with pytest.raises(ValueError, match = "path components"):
export_models._validate_save_directory(str(tmp_path / ("a" * 256)))
def test_export_write_dir_accepts_external_absolute_but_read_dir_rejects(tmp_path, monkeypatch):
storage_roots = _load_module(
"test_storage_roots_accept_external",
"utils/paths/storage_roots.py",
)
export_root = tmp_path / "exports"
external = tmp_path / "external"
export_root.mkdir()
external.mkdir()
monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root)
assert storage_roots.resolve_export_write_dir(str(external)) == external
with pytest.raises(ValueError, match = "path escapes root"):
storage_roots.resolve_export_dir(str(external))
def test_export_write_dir_accepts_expanded_home_path(tmp_path, monkeypatch):
storage_roots = _load_module(
"test_storage_roots_accept_home_path",
"utils/paths/storage_roots.py",
)
export_root = tmp_path / "exports"
home = tmp_path / "home"
export_root.mkdir()
home.mkdir()
monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root)
if storage_roots.os.name == "nt":
monkeypatch.setenv("USERPROFILE", str(home))
else:
monkeypatch.setenv("HOME", str(home))
assert storage_roots.resolve_export_write_dir("~/exports/model") == home / "exports" / "model"
def test_resolve_export_write_dir_rejects_backslash_parent_segment():
storage_roots = _load_module(
"test_storage_roots_reject_parent",
"utils/paths/storage_roots.py",
)
with pytest.raises(ValueError, match = r"\.\."):
storage_roots.resolve_export_write_dir(r"exports\..\outside")
def test_export_write_dir_handles_non_native_windows_absolute_as_relative(tmp_path, monkeypatch):
storage_roots = _load_module(
"test_storage_roots_non_native_windows_path",
"utils/paths/storage_roots.py",
)
export_root = tmp_path / "exports"
export_root.mkdir()
monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root)
if storage_roots.os.name == "nt":
pytest.skip("Windows drive paths are native on Windows")
assert (
storage_roots.resolve_export_write_dir(r"C:\exports\model")
== export_root / r"C:\exports\model"
)
def test_export_details_registers_external_absolute_output(tmp_path, monkeypatch):
_install_lightweight_backend_stubs(monkeypatch)
export_route = _load_module(
"test_routes_export_external",
"routes/export.py",
monkeypatch,
)
output = tmp_path / "Gemma4_26B_gguf"
output.mkdir()
export_root = tmp_path / "studio" / "exports"
export_root.mkdir(parents = True)
registered = []
monkeypatch.setattr(
export_route,
"_try_register_external_export",
lambda path: (registered.append(path) is None, str(path)),
)
monkeypatch.setattr(
"utils.paths.storage_roots.exports_root",
lambda: export_root,
)
details = export_route._export_details(str(output))
assert details == {
"output_path": str(output),
"scan_folder_registered": True,
"scan_folder_path": str(output),
}
assert registered == [output]
def test_export_details_does_not_register_contained_exports(tmp_path, monkeypatch):
_install_lightweight_backend_stubs(monkeypatch)
export_route = _load_module(
"test_routes_export_contained",
"routes/export.py",
monkeypatch,
)
export_root = tmp_path / "exports"
output = export_root / "model-gguf"
output.mkdir(parents = True)
monkeypatch.setattr(
export_route,
"_try_register_external_export",
lambda path: pytest.fail(f"unexpected registration: {path}"),
)
monkeypatch.setattr(
"utils.paths.storage_roots.exports_root",
lambda: export_root,
)
assert export_route._export_details(str(output)) == {"output_path": "model-gguf"}
def test_registered_absolute_export_folder_is_discoverable(tmp_path, monkeypatch):
_install_lightweight_backend_stubs(monkeypatch)
models_route = _load_module("test_routes_models", "routes/models.py", monkeypatch)
export_dir = tmp_path / "Gemma4_26B_gguf"
export_dir.mkdir()
gguf_file = export_dir / "Gemma4_26B.BF16-00001-of-00002.gguf"
gguf_file.write_bytes(b"gguf")
found = models_route._scan_models_dir(export_dir)
assert len(found) == 1
assert found[0].path == str(gguf_file)
assert found[0].source == "models_dir"

139
studio/backend/tests/test_external_provider_usage_chunk.py
View file

@ -144,6 +144,14 @@ def _make_openai_client() -> ExternalProviderClient:
)
def _make_custom_client() -> ExternalProviderClient:
return ExternalProviderClient(
provider_type = "custom",
base_url = "http://custom.example/v1",
api_key = "",
)
def _anthropic_sse(events: list[dict]) -> bytes:
chunks: list[str] = []
for event in events:
@ -180,6 +188,137 @@ def _usage_chunks(lines: list[str]) -> list[dict]:
return out
def test_custom_provider_registry_is_hidden():
from core.inference.providers import get_provider_info, list_available_providers
info = get_provider_info("custom")
assert info is not None
assert info["hidden"] is True
assert "custom" not in {p["provider_type"] for p in list_available_providers()}
def test_custom_provider_uses_chat_completions_without_auth_key(monkeypatch):
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["url"] = str(request.url)
captured["headers"] = dict(request.headers)
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\ndata: [DONE]\n\n',
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_custom_client()
lines = await _collect(
client.stream_chat_completion(
messages = [{"role": "user", "content": "ping"}],
model = "Qwen/Qwen3-0.6B",
temperature = 0.7,
top_p = 0.95,
max_tokens = 64,
)
)
await client.close()
return lines
lines = _drive(run())
assert captured["url"] == "http://custom.example/v1/chat/completions"
assert "authorization" not in {k.lower() for k in captured["headers"]}
assert captured["body"]["model"] == "Qwen/Qwen3-0.6B"
assert any("ok" in line for line in lines)
def test_custom_provider_test_endpoint_probes_chat_completion(monkeypatch):
import importlib.util
import sys
from pathlib import Path
module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py"
spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path)
assert spec is not None
assert spec.loader is not None
providers_route = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = providers_route
spec.loader.exec_module(providers_route)
captured: dict = {}
class _FakeClient:
def __init__(self, **kwargs):
captured["init"] = kwargs
async def chat_completion(self, **kwargs):
captured["chat_completion"] = kwargs
return {"choices": [{"message": {"content": "ok"}}]}
async def list_models(self):
raise AssertionError("custom provider test must not call /models")
async def close(self):
captured["closed"] = True
monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient)
async def run():
return await providers_route.test_provider(
providers_route.ProviderTestRequest(
provider_type = "custom",
base_url = "http://custom.example/v1",
model_id = "Qwen/Qwen3-0.6B",
),
current_subject = "unsloth",
)
result = _drive(run())
assert result.success is True
assert result.models_count is None
assert captured["init"]["provider_type"] == "custom"
assert captured["chat_completion"]["model"] == "Qwen/Qwen3-0.6B"
assert captured["chat_completion"]["max_tokens"] == 1
assert captured["closed"] is True
def test_custom_provider_test_endpoint_requires_model_id(monkeypatch):
import importlib.util
import sys
from pathlib import Path
module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py"
spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path)
assert spec is not None
assert spec.loader is not None
providers_route = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = providers_route
spec.loader.exec_module(providers_route)
class _FakeClient:
def __init__(self, **kwargs):
pass
async def close(self):
pass
monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient)
async def run():
return await providers_route.test_provider(
providers_route.ProviderTestRequest(
provider_type = "custom",
base_url = "http://custom.example/v1",
),
current_subject = "unsloth",
)
result = _drive(run())
assert result.success is False
assert "model ID" in result.message
def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch):
sse_events = [
{

345
studio/backend/tests/test_gemma4_chat_template_override.py Normal file
View file

@ -0,0 +1,345 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Auto-override of the chat template for ``unsloth/gemma-4-*-GGUF``.
Studio ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking``
defaulted off) and applies it to gemma-4 GGUF loads via the existing
``chat_template_override`` -> ``--chat-template-file`` path, so users do not need
to re-download quants. Pins the family matcher, the resolver precedence, the
bundled asset's reasoning/tool capabilities (which drive the "Preserve thinking"
UI toggle), the Jinja gate behaviour, and the reload-dedup interaction.
"""
from __future__ import annotations
import importlib.util
import sys
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import pytest
# ── chat_templates is dependency-light: load it directly so the pure-logic
# tests run without the studio venv / core.inference package side effects. ──
_CT_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_templates.py"
_ct_spec = importlib.util.spec_from_file_location("_gemma4_ct_test", _CT_PATH)
chat_templates = importlib.util.module_from_spec(_ct_spec)
_ct_spec.loader.exec_module(chat_templates)
is_unsloth_gemma4_gguf = chat_templates.is_unsloth_gemma4_gguf
resolve_effective_chat_template_override = chat_templates.resolve_effective_chat_template_override
load_bundled_chat_template = chat_templates.load_bundled_chat_template
is_unsloth_gemma4_edge_gguf = chat_templates.is_unsloth_gemma4_edge_gguf
BUNDLED = load_bundled_chat_template("gemma-4.jinja") # 12b / 26B-A4B / 31B
EDGE = load_bundled_chat_template("gemma-4-edge.jinja") # E2B / E4B
# ── Stubs so core.inference.llama_cpp imports without the full studio venv ──
def _stub_modules_ctx():
"""patch.dict context that stubs the heavy deps llama_cpp pulls in at import,
but only those NOT already importable (real httpx / structlog are kept when
present, e.g. in CI), and removes the stubs on exit so other tests are not
polluted."""
from unittest.mock import patch
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
_httpx_stub.Client = type(
"C",
(),
{
"__init__": lambda s, **kw: None,
"__enter__": lambda s: s,
"__exit__": lambda s, *a: None,
},
)
overrides = {
name: stub
for name, stub in (
("loggers", _loggers_stub),
("structlog", _structlog_stub),
("httpx", _httpx_stub),
)
if name not in sys.modules
}
return patch.dict(sys.modules, overrides)
def _detect_reasoning_flags():
with _stub_modules_ctx():
from core.inference.llama_cpp import detect_reasoning_flags
return detect_reasoning_flags
# ── Family matcher ───────────────────────────────────────────────────
@pytest.mark.parametrize(
"model_id,expected",
[
("unsloth/gemma-4-E2B-it-GGUF", True),
("unsloth/gemma-4-E4B-it-GGUF", True),
("unsloth/gemma-4-31B-it-GGUF", True),
("unsloth/gemma-4-26B-A4B-it-GGUF", True),
("UNSLOTH/GEMMA-4-E2B-IT-GGUF", True), # case-insensitive
("gemma-4-E2B-it-GGUF", True), # owner-less shorthand -> unsloth/
("gemma-4-31B-it-GGUF", True), # owner-less shorthand -> unsloth/
("unsloth/gemma-4-E2B-it", False), # bf16, not GGUF
("unsloth/gemma-3-4b-it-GGUF", False), # gemma 3
("google/gemma-4-31B-it-GGUF", False), # not unsloth
("unsloth/Qwen3.5-9B-MTP-GGUF", False),
("/home/user/models/gemma-4-E2B.Q4_K_M.gguf", False), # local path
("", False),
(None, False),
],
)
def test_is_unsloth_gemma4_gguf(model_id, expected):
assert is_unsloth_gemma4_gguf(model_id) is expected
# ── Resolver precedence ──────────────────────────────────────────────
@pytest.mark.parametrize(
"model_id,expected_edge",
[
("unsloth/gemma-4-E2B-it-GGUF", True),
("unsloth/gemma-4-E4B-it-GGUF", True),
("UNSLOTH/GEMMA-4-E4B-IT-GGUF", True),
("unsloth/gemma-4-12b-it-GGUF", False),
("unsloth/gemma-4-26B-A4B-it-GGUF", False),
("unsloth/gemma-4-31B-it-GGUF", False),
("unsloth/gemma-3-4b-it-GGUF", False),
],
)
def test_is_unsloth_gemma4_edge_gguf(model_id, expected_edge):
assert is_unsloth_gemma4_edge_gguf(model_id) is expected_edge
def test_resolver_returns_edge_template_for_e2b_e4b():
for mid in ("unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF"):
out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None)
assert out == EDGE
assert out != BUNDLED
def test_resolver_handles_owner_less_shorthand():
# ModelConfig.from_identifier prefixes unsloth/ for bare ids; the resolver
# runs before that, so it must apply the same normalization.
assert (
resolve_effective_chat_template_override(
model_identifier = "gemma-4-E2B-it-GGUF", user_override = None
)
== EDGE
)
assert (
resolve_effective_chat_template_override(
model_identifier = "gemma-4-31B-it-GGUF", user_override = None
)
== BUNDLED
)
def test_resolver_returns_standard_template_for_larger_models():
for mid in (
"unsloth/gemma-4-12b-it-GGUF",
"unsloth/gemma-4-26B-A4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",
):
out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None)
assert out == BUNDLED
def test_resolver_user_override_wins():
out = resolve_effective_chat_template_override(
model_identifier = "unsloth/gemma-4-E2B-it-GGUF", user_override = "MY TEMPLATE"
)
assert out == "MY TEMPLATE"
def test_resolver_blank_override_falls_back_to_bundled():
out = resolve_effective_chat_template_override(
model_identifier = "unsloth/gemma-4-31B-it-GGUF", user_override = " "
)
assert out == BUNDLED
def test_resolver_none_for_non_gemma():
assert (
resolve_effective_chat_template_override(
model_identifier = "unsloth/Llama-3.2-1B-Instruct-GGUF", user_override = None
)
is None
)
# ── Bundled asset content + capability classification ────────────────
@pytest.mark.parametrize("tpl", [BUNDLED, EDGE])
def test_bundled_template_has_preserve_thinking_defaulted_off(tpl):
assert "preserve_thinking" in tpl
assert "preserve_thinking | default(false)" in tpl
@pytest.mark.parametrize("name", ["gemma-4.jinja", "gemma-4-edge.jinja"])
def test_bundled_templates_are_ascii(name):
# The temp file written for --chat-template-file must encode on any locale.
# Keeping the bundled templates ASCII avoids UnicodeEncodeError on non-UTF-8
# Windows locales (cp932/cp1252) regardless of the writer's encoding.
text = load_bundled_chat_template(name)
non_ascii = sorted({c for c in text if ord(c) > 127})
assert not non_ascii, f"{name} has non-ASCII chars: {non_ascii}"
@pytest.mark.parametrize("tpl", [BUNDLED, EDGE])
def test_detect_reasoning_flags_on_bundled_template(tpl):
detect_reasoning_flags = _detect_reasoning_flags()
flags = detect_reasoning_flags(tpl, "unsloth/gemma-4-E2B-it-GGUF")
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking"
assert flags["reasoning_always_on"] is False
# This is what makes the "Preserve thinking" toggle appear in the UI.
assert flags["supports_preserve_thinking"] is True
assert flags["supports_tools"] is True
def test_edge_template_omits_empty_thought_block_on_thinking_off():
"""E2B/E4B must NOT emit the empty <|channel>thought<channel|> block when
thinking is disabled; the larger-model template must. This is the only
intended difference between the two bundled templates."""
EMPTY = "<|channel>thought\n<channel|>"
msgs = [{"role": "user", "content": "hi"}]
edge_off = _render_with(EDGE, msgs, enable_thinking = False)
std_off = _render_with(BUNDLED, msgs, enable_thinking = False)
assert EMPTY not in edge_off, "edge (E2B/E4B) should not emit empty thought block"
assert EMPTY in std_off, "standard (12b/26B/31B) should emit empty thought block"
# With thinking ON neither appends the empty block at the prompt tail.
assert EMPTY not in _render_with(EDGE, msgs, enable_thinking = True)
# ── Jinja gate behaviour (off = omit prior reasoning, on = keep) ─────
def _render_with(tpl, messages, **kw):
pytest.importorskip("jinja2") # transitive via transformers; skip in minimal envs
from jinja2 import Environment, BaseLoader
def raise_exception(msg):
raise RuntimeError(msg)
env = Environment(loader = BaseLoader())
return env.from_string(tpl).render(
messages = messages,
bos_token = "<bos>",
raise_exception = raise_exception,
add_generation_prompt = True,
**kw,
)
def _render(messages, **kw):
return _render_with(BUNDLED, messages, **kw)
def _convo_with_prior_tool_reasoning():
# Assistant tool-call turn with reasoning, BEFORE the last user message.
return [
{"role": "user", "content": "q1"},
{
"role": "assistant",
"reasoning_content": "SECRET_THOUGHT",
"tool_calls": [{"id": "c1", "function": {"name": "f", "arguments": {"x": 1}}}],
},
{"role": "tool", "tool_call_id": "c1", "content": "42"},
{"role": "user", "content": "q2"},
]
def test_preserve_thinking_off_omits_prior_reasoning():
# default(false): kwarg unset -> prior reasoning dropped before last user turn.
assert "SECRET_THOUGHT" not in _render(_convo_with_prior_tool_reasoning())
def test_preserve_thinking_on_keeps_prior_reasoning():
assert "SECRET_THOUGHT" in _render(_convo_with_prior_tool_reasoning(), preserve_thinking = True)
def test_enable_thinking_gates_think_token():
assert "<|think|>" in _render([{"role": "user", "content": "hi"}], enable_thinking = True)
assert "<|think|>" not in _render([{"role": "user", "content": "hi"}], enable_thinking = False)
# ── Reload dedup interaction (why the route resolves the effective override) ──
def test_already_in_target_state_consistent_with_bundled_override():
"""The backend dedup compares the incoming override against the live one.
The route resolves the bundled template up front so a re-load that omits
``chat_template_override`` still matches (no spurious reload), while a raw
``None`` would not.
"""
LlamaCppBackend = _import_backend()
class _FakeProcess:
def terminate(self): ...
def wait(self, timeout = None):
return 0
def kill(self): ...
def poll(self):
return 0
backend = LlamaCppBackend()
backend._process = _FakeProcess()
backend._healthy = True
backend._model_identifier = "unsloth/gemma-4-E2B-it-GGUF"
backend._hf_variant = "Q4_K_M"
backend._requested_n_ctx = 8192
backend._cache_type_kv = None
backend._speculative_type = None
backend._requested_spec_mode = "auto"
backend._chat_template_override = BUNDLED # live server launched with the bundle
backend._is_vision = False
backend._extra_args = None
backend._gguf_path = None
common = dict(
model_identifier = "unsloth/gemma-4-E2B-it-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
extra_args = None,
is_vision = False,
)
# Effective (resolved bundled) override -> already loaded, no reload.
assert backend._already_in_target_state(chat_template_override = BUNDLED, **common) is True
# Raw None (unresolved) -> false match, would force a needless reload.
assert backend._already_in_target_state(chat_template_override = None, **common) is False
def _import_backend():
with _stub_modules_ctx():
from core.inference.llama_cpp import LlamaCppBackend
return LlamaCppBackend

7
studio/backend/tests/test_gguf_route_cursor_reset.py
View file

@ -59,11 +59,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
payload,
_cancel_event,
headers = None,
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
def fake_iter_text_cancellable(response, _cancel_event):
def fake_iter_text_cancellable(
response,
_cancel_event,
first_token_deadline = None,
):
yield from response.chunks
monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)

523
studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
View file

@ -1,523 +0,0 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets.
Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade
GitHub API are stubbed out so the suite runs without internet access and is
not subject to rate limits.
"""
from __future__ import annotations
import importlib
import sys
from pathlib import Path
from unittest.mock import patch
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
_mod = importlib.import_module("install_llama_prebuilt")
HostInfo = _mod.HostInfo
resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None)
_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None)
if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
pytest.skip("PR symbols not present - check branch", allow_module_level = True)
@pytest.fixture(autouse = True)
def _clear_lemonade_release_cache():
"""Prevent cross-test pollution of the lemonade release lru_cache and
selection-log dedup set when tests vary the fetch_json mock return value."""
_cache = getattr(_mod, "_fetch_lemonade_release_cached", None)
_logged: set | None = getattr(_mod, "_lemonade_selection_logged", None)
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
if _logged is not None:
_logged.clear()
yield
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
if _logged is not None:
_logged.clear()
_STUB_TAG = "b1262"
_STUB_OS_PREFIXES = ("ubuntu", "windows")
_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X")
def _stub_lemonade_release() -> dict:
"""Minimal lemonade release payload covering all supported GPU/OS combinations."""
assets = [
{
"name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip",
"browser_download_url": (
f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/"
f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip"
),
}
for prefix in _STUB_OS_PREFIXES
for family in _STUB_FAMILIES
]
return {"tag_name": _STUB_TAG, "assets": assets}
def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo:
return HostInfo(
system = "Windows" if windows else "Linux",
machine = "amd64" if windows else "x86_64",
is_windows = windows,
is_linux = not windows,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = True,
rocm_gfx_target = gfx_target,
)
def _lookup_family(gfx: str) -> str | None:
for prefix, family in _LEMONADE_GFX_FAMILIES:
if gfx.startswith(prefix):
return family
return None
# ---------------------------------------------------------------------------
# GPU family mapping
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"gfx,expected_family",
[
("gfx1151", "gfx1151"),
("gfx1150", "gfx1150"),
("gfx1201", "gfx120X"),
("gfx1200", "gfx120X"),
("gfx1100", "gfx110X"),
("gfx1030", "gfx103X"),
],
)
def test_gpu_family_mapping(gfx, expected_family):
assert _lookup_family(gfx) == expected_family
def test_unknown_gpu_not_in_families():
assert _lookup_family("gfx999") is None
# ---------------------------------------------------------------------------
# Asset resolution - hits real lemonade GitHub API
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"gfx,os_prefix,windows",
[
("gfx1151", "ubuntu", False),
("gfx1150", "ubuntu", False),
("gfx1201", "ubuntu", False),
("gfx1100", "ubuntu", False),
("gfx1030", "ubuntu", False),
("gfx1151", "windows", True),
("gfx1100", "windows", True),
],
)
def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows):
host = _make_rocm_host(gfx, windows = windows)
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
result = resolve_lemonade_rocm_choice(host, os_prefix, "default", llama_tag = "latest")
assert result is not None, f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})"
assert _lookup_family(gfx) in result.name
assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm")
def test_unknown_gpu_falls_through_to_upstream():
host = _make_rocm_host("gfx999")
result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest")
assert result is None
# ---------------------------------------------------------------------------
# The Linux attempt builder must plan a lemonade ROCm attempt for AMD-only hosts.
# This is the path setup.sh actually invokes (fork hosts now select from the
# manifest), so the lemonade integration is useless if it isn't wired in here.
# ---------------------------------------------------------------------------
_linux_published_attempts = getattr(_mod, "_linux_published_attempts", None)
direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None)
PublishedLlamaArtifact = _mod.PublishedLlamaArtifact
PublishedReleaseBundle = _mod.PublishedReleaseBundle
def _rocm_bundle(gfx_family: str, mapped_targets: list[str]) -> "PublishedReleaseBundle":
"""A fork manifest bundle exposing a per-gfx linux-rocm artifact, so
published_rocm_choice_for_host can match the host before the lemonade
fallback is appended."""
asset_name = f"app-b9457-linux-x64-rocm-{gfx_family}.tar.gz"
artifact = PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = "linux-rocm",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
rank = 1000,
gfx_target = gfx_family,
mapped_targets = mapped_targets,
)
return PublishedReleaseBundle(
repo = "unslothai/llama.cpp",
release_tag = "v1.0",
upstream_tag = "b9457",
assets = {asset_name: f"https://example.invalid/{asset_name}"},
artifacts = [artifact],
)
@pytest.mark.skipif(
_linux_published_attempts is None,
reason = "Linux attempt builder not present on this branch",
)
def test_linux_attempts_include_fork_rocm_and_lemonade_for_rocm_host():
host = _make_rocm_host("gfx1151")
bundle = _rocm_bundle("gfx1151", ["gfx1151"])
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
attempts = _linux_published_attempts(host, bundle, "latest")
kinds = [a.install_kind for a in attempts]
assert "linux-rocm" in kinds, f"builder did not include any linux-rocm attempt; got {kinds}"
sources = {a.source_label for a in attempts if a.install_kind == "linux-rocm"}
# The fork's own per-gfx bundle is preferred, with the lemonade prebuilt as
# the fallback -- both must be present for a covered ROCm host.
assert "published" in sources, f"fork ROCm bundle missing; got {sources}"
assert "lemonade" in sources, f"lemonade ROCm fallback missing; got {sources}"
lemonade_attempt = next(a for a in attempts if a.source_label == "lemonade")
assert "gfx1151" in lemonade_attempt.name
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "direct release planners not present on this branch",
)
def test_direct_upstream_plan_includes_lemonade_for_windows_hip_host():
host = _make_rocm_host("gfx1151", windows = True)
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [],
}
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt"
kinds = [a.install_kind for a in plan.attempts]
assert "windows-hip" in kinds, f"planner did not include a lemonade HIP attempt; got {kinds}"
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "direct release planners not present on this branch",
)
def test_windows_hip_falls_back_to_upstream_when_lemonade_unavailable():
"""If lemonade returns None (e.g. gfx999 or transient API failure), the planner
must still include the upstream HIP asset rather than silently downgrading to CPU."""
host = _make_rocm_host("gfx999", windows = True)
hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip"
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [
{
"name": hip_asset,
"browser_download_url": f"https://example.invalid/{hip_asset}",
},
],
}
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None
kinds = [a.install_kind for a in plan.attempts]
assert "windows-hip" in kinds, f"upstream HIP asset not included as fallback; got {kinds}"
hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip")
assert hip_attempt.source_label == "upstream"
# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ──
def test_lemonade_release_api_url_pinned_tag():
"""A pinned llama_tag must produce the /releases/tags/<tag> URL."""
assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262")
assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest")
assert _mod._lemonade_release_api_for("").endswith("/releases/latest")
def test_lemonade_release_api_url_encodes_tag():
"""Unexpected slashes / hashes in the tag must be URL-encoded so the URL
cannot be reshaped (defence in depth -- tags should already be sanitised
upstream)."""
url = _mod._lemonade_release_api_for("b1260/../latest")
assert "/releases/tags/b1260%2F..%2Flatest" in url
assert "//latest" not in url.split("/releases/tags/", 1)[1]
def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch):
"""UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver."""
monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1")
host = _make_rocm_host("gfx1151")
res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
"""If the GitHub API response somehow contained an off-host download URL,
the resolver must refuse to use it (lemonade assets are not in the
approved-hash manifest)."""
bad_release = {
"tag_name": _STUB_TAG,
"assets": [
{
"name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
"browser_download_url": "https://attacker.invalid/llama.zip",
},
],
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = bad_release):
res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
def test_lemonade_resolver_rejects_http_scheme():
assert not _mod._is_trusted_github_release_url(
"http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_accepts_github_cdn():
# Real GitHub release CDN URLs carry the /github-production-release-asset- prefix.
assert _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_rejects_arbitrary_cdn_path():
# A CDN URL without the release-asset path prefix must be rejected.
assert not _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/abc/def",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_accepts_release_path():
url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip"
assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm")
def test_lemonade_resolver_rejects_wrong_repo():
"""A github.com release URL for a different repo must be rejected."""
assert not _mod._is_trusted_github_release_url(
"https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_rejects_empty_browser_download_url():
"""An asset entry with an empty browser_download_url must fall through."""
release = {
"tag_name": _STUB_TAG,
"assets": [
{
"name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
"browser_download_url": "",
},
],
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = release):
res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
def test_lemonade_runtime_patterns_include_hip_runtime():
"""linux-rocm overlay must use a broad lib glob to catch all bundled .so files.
Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp,
...) whose names change across ROCm releases. A broad ``lib*.so*`` glob
avoids having to enumerate every transitive dependency by name.
"""
from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice
choice = AssetChoice(
repo = "lemonade-sdk/llamacpp-rocm",
tag = "b1262",
name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip",
url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip",
source_label = "lemonade",
install_kind = "linux-rocm",
)
pats = runtime_patterns_for_choice(choice)
# The broad glob must be present so every .so in the lemonade bundle
# (including transitive deps added in future ROCm releases) gets overlaid.
assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}"
_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None)
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch):
"""AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES;
on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
# Two GPUs; rocminfo reports each token twice (as in the real tool output).
probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
assert _pick_rocm_gfx_target(probe_out) == "gfx1100"
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch):
"""CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None."""
probe_out = "gfx1151\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1")
assert _pick_rocm_gfx_target(probe_out) is None
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
"""Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must
return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the
two gfx1100 entries into one and making index 2 out of range."""
# Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
# Each GPU gets its own Agent section with a few token mentions.
probe_out = (
"***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n"
"***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n"
"***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n"
)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2")
assert _pick_rocm_gfx_target(probe_out) == "gfx1151"
# ---------------------------------------------------------------------------
# Fork release scan: Windows ROCm resolves lemonade by the requested tag
# ---------------------------------------------------------------------------
_resolve_release_asset_choice = getattr(_mod, "resolve_release_asset_choice", None)
_ApprovedReleaseChecksums = getattr(_mod, "ApprovedReleaseChecksums", None)
@pytest.mark.skipif(
_resolve_release_asset_choice is None or _ApprovedReleaseChecksums is None,
reason = "fork release planner not present on this branch",
)
def test_fork_scan_windows_rocm_resolves_lemonade_by_requested_tag():
"""The fork release scan pins llama_tag to per-release upstream tags
(b9457, ...) that lemonade's own tag series never contains, so the
lemonade lookup must use the requested tag ("latest") instead. Pinning
lemonade to the per-release tag 404s on every scanned release and a
Windows ROCm host ends in a rate-limited fatal instead of the lemonade
prebuilt."""
host = _make_rocm_host("gfx1151", windows = True)
# No windows-rocm artifact in the bundle, matching current fork releases.
bundle = _rocm_bundle("gfx1151", ["gfx1151"])
checksums = _ApprovedReleaseChecksums(
repo = "unslothai/llama.cpp",
release_tag = "v1.0",
upstream_tag = "b9457",
artifacts = {},
)
seen_urls: list[str] = []
def _fake_fetch(api_url, *args, **kwargs):
seen_urls.append(api_url)
if "lemonade-sdk" in api_url:
if api_url.endswith("/releases/latest"):
return _stub_lemonade_release()
raise RuntimeError(f"unexpected pinned lemonade fetch: {api_url}")
# ggml-org asset listing for the upstream HIP/CPU filename fallbacks.
return {"tag_name": "b9457", "assets": []}
with patch.object(_mod, "fetch_json", side_effect = _fake_fetch):
attempts = _resolve_release_asset_choice(
host,
"b9457", # concrete per-release upstream tag from the scan loop
bundle,
checksums,
requested_tag = "latest",
)
lemonade = [a for a in attempts if a.source_label == "lemonade"]
assert lemonade, f"lemonade attempt missing for Windows ROCm host; got {attempts}"
assert "gfx1151" in lemonade[0].name
assert any(
u.endswith("/releases/latest") for u in seen_urls
), f"lemonade was never resolved via /releases/latest; fetches: {seen_urls}"
assert not any(
"lemonade-sdk" in u and "/releases/tags/" in u for u in seen_urls
), f"lemonade lookup was pinned to the fork release tag: {seen_urls}"
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "direct release planners not present on this branch",
)
def test_direct_upstream_plan_includes_lemonade_for_linux_rocm_host():
"""A Linux ROCm host on the ggml-org direct path (e.g. a --published-repo
override) must plan lemonade before the CPU tarball, mirroring the Windows
branch. The lemonade planning previously lived in the removed
--simple-policy dispatcher, so without this leg such hosts silently
install the CPU build."""
host = _make_rocm_host("gfx1151")
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [
{
"name": "llama-b9022-bin-ubuntu-x64.tar.gz",
"browser_download_url": (
"https://github.com/ggml-org/llama.cpp/releases/download/"
"b9022/llama-b9022-bin-ubuntu-x64.tar.gz"
),
}
],
}
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None, "Linux ROCm host should produce a direct plan"
kinds = [a.install_kind for a in plan.attempts]
sources = [a.source_label for a in plan.attempts]
assert "linux-rocm" in kinds, f"lemonade ROCm attempt missing; got {kinds}"
assert sources[0] == "lemonade", f"lemonade must be the first attempt; got {sources}"
assert "gfx1151" in plan.attempts[0].name

87
studio/backend/tests/test_llama_cpp_freshness.py
View file

@ -433,3 +433,90 @@ def test_fetch_latest_release_tag_uses_publish_time(monkeypatch):
]
monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload))
assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453"
# reset_caches(drop_disk=...) -- post-update stale same-base mix disk cache.
def _seed_disk_cache(tmp_path: Path, latest_tag: str) -> Path:
# Matches _cache_path_for under the fixture's stubbed _cache_dir.
cache_dir = tmp_path / ".freshness"
cache_dir.mkdir(exist_ok = True)
cache_file = cache_dir / "unslothai__llama.cpp.json"
cache_file.write_text(json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag}))
return cache_file
def test_reset_caches_drop_disk_removes_disk_cache(tmp_path):
cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa")
assert cache_file.exists()
fr.reset_caches(drop_disk = True)
assert not cache_file.exists()
def test_reset_caches_default_keeps_disk_cache(tmp_path):
# The no-arg form is in-memory only (its existing test-only contract); it
# must not delete the on-disk cache.
cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa")
fr.reset_caches()
assert cache_file.exists()
def test_reset_caches_drop_disk_on_missing_dir_is_noop(tmp_path):
# Fresh machine, no cache dir yet: drop_disk must be a quiet no-op.
assert not (tmp_path / ".freshness").exists()
fr.reset_caches(drop_disk = True) # must not raise
def test_drop_disk_lets_banner_fail_open_after_same_base_mix_swap(monkeypatch, tmp_path):
# P2 #2: the disk cache holds a still-fresh same-base mix (b9596-mix-aaa)
# from before an update to a *different* same-base mix (b9596-mix-bbb).
# The post-install path drops the disk cache; if the forced refresh is then
# offline, latest reads as None and the banner fails open -- instead of
# replaying the stale b9596-mix-aaa and falsely reading "behind".
_seed_disk_cache(tmp_path, "b9596-mix-aaa")
install_dir = tmp_path / "llama.cpp"
_write_marker(
install_dir,
tag = "b9596",
release_tag = "b9596-mix-bbb",
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5))
.isoformat()
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
# GitHub unreachable for the rest of the test (the offline post-install
# refresh, and the later status check).
monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
fr.reset_caches(drop_disk = True) # exactly what the apply path now does
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["latest_tag"] is None
assert info["behind"] is False
assert info["stale"] is False
def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path):
# Contrast/guard for the case above: an in-memory-only reset leaves the
# stale same-base mix on disk, so an offline check replays it and falsely
# reads behind/stale. This is exactly the failure drop_disk removes; if a
# future change makes the no-arg reset also clear disk, the apply-path call
# and this guard should be revisited together.
_seed_disk_cache(tmp_path, "b9596-mix-aaa")
install_dir = tmp_path / "llama.cpp"
_write_marker(
install_dir,
tag = "b9596",
release_tag = "b9596-mix-bbb",
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5))
.isoformat()
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
fr.reset_caches() # in-memory only -> stale disk value survives
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["latest_tag"] == "b9596-mix-aaa"
assert info["behind"] is True
assert info["stale"] is True

178
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -21,6 +21,8 @@ if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference.llama_cpp import LlamaCppBackend
from state import tool_approvals
from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision
def _sse(delta: dict) -> str:
@ -50,11 +52,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
payload,
_cancel_event,
headers = None,
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
def fake_iter_text_cancellable(response, _cancel_event):
def fake_iter_text_cancellable(
response,
_cancel_event,
first_token_deadline = None,
):
yield from response.chunks
monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)
@ -70,6 +77,27 @@ def _tool_names(payload: dict) -> list[str]:
]
def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list[str]:
return [
_sse(
{
"tool_calls": [
{
"index": 0,
"id": call_id,
"type": "function",
"function": {
"name": tool_name,
"arguments": json.dumps(arguments),
},
}
]
}
),
_done(),
]
def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
"""llama-server may emit content first and then native delta.tool_calls.
@ -1149,3 +1177,151 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
assert content_texts == ["I will use render_html now.", "Final note after tool."]
assert len(payloads) == 3
def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch):
streams = [
_structured_tool_call("python", {"code": "print(1)"}, "call_py"),
[_sse({"content": "Done."}), _done()],
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, streams, payloads)
calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "OK"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: "approval-1")
monkeypatch.setattr(
"core.inference.llama_cpp.begin_tool_decision",
lambda *_a, **_k: object(),
)
monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow")
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run python"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
confirm_tool_calls = True,
session_id = "sess",
)
)
starts = [event for event in events if event.get("type") == "tool_start"]
assert len(starts) == 1
assert starts[0]["approval_id"]
assert starts[0]["awaiting_confirmation"] is True
assert calls == [("python", {"code": "print(1)"})]
assert any(event.get("type") == "tool_end" and event.get("result") == "OK" for event in events)
def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch):
approval_id = "approval-close"
streams = [_structured_tool_call("python", {"code": "print(1)"}, "call_py")]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, streams, payloads)
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("tool should not run")),
)
monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: approval_id)
with tool_approvals._lock:
tool_approvals._pending.clear()
gen = backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run python"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
confirm_tool_calls = True,
session_id = "sess",
)
try:
assert next(gen)["type"] == "status"
start = next(gen)
assert start["type"] == "tool_start"
assert start["approval_id"] == approval_id
with tool_approvals._lock:
assert approval_id in tool_approvals._pending
finally:
gen.close()
with tool_approvals._lock:
assert approval_id not in tool_approvals._pending
assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False
def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch):
streams = [[_sse({"content": "Done."}), _done()]]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, streams, payloads)
def fail_autoinject(*_args, **_kwargs):
raise AssertionError("RAG autoinject must not run before approval")
monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "use docs"}],
tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}],
max_tool_iterations = 1,
confirm_tool_calls = True,
session_id = "sess",
rag_scope = {"thread_id": "t1"},
)
)
assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events)
def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypatch):
same_call = _structured_tool_call("python", {"code": "print(1)"}, "call_py")
streams = [
same_call,
_structured_tool_call("python", {"code": "print(1)"}, "call_py_retry"),
[_sse({"content": "Done."}), _done()],
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, streams, payloads)
calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "OK"
decisions = iter(["deny", "allow"])
approvals = iter(["approval-1", "approval-2"])
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: next(approvals))
monkeypatch.setattr(
"core.inference.llama_cpp.begin_tool_decision",
lambda *_a, **_k: object(),
)
monkeypatch.setattr(
"core.inference.llama_cpp.wait_tool_decision",
lambda *_a, **_k: next(decisions),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run python"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 2,
confirm_tool_calls = True,
session_id = "sess",
)
)
starts = [event for event in events if event.get("type") == "tool_start"]
ends = [event for event in events if event.get("type") == "tool_end"]
assert len(starts) == 2
assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"]
assert calls == [("python", {"code": "print(1)"})]

37
studio/backend/tests/test_llama_cpp_update.py
View file

@ -233,6 +233,39 @@ def test_status_source_build_suppressed_when_newer(monkeypatch, tmp_path):
assert st["installed_tag"] == "b9600"
def test_status_source_build_offers_same_base_mix(monkeypatch, tmp_path):
# The reported banner bug: a source build at the same upstream base as a new
# Unsloth prebuilt that adds a mix-<sha> suffix. The base build numbers match
# (9596 == 9596) but the mix carries extra patches the source build lacks, so
# the update must still surface -- mirroring the marker path's is_behind.
binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(monkeypatch, release_tag = "b9596-mix-e6f2453", llama_tag = "b9596")
monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9596)
st = upd.get_update_status()
assert st["supported"] is True
assert st["update_available"] is True
assert st["source_build"] is True
assert st["installed_tag"] == "b9596"
assert st["latest_tag"] == "b9596-mix-e6f2453"
def test_status_source_build_same_base_bare_not_offered(monkeypatch, tmp_path):
# Same base, but the prebuilt is a bare rebuild (no mix suffix): nothing extra
# to gain, so do not nag.
binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
binary.parent.mkdir(parents = True)
binary.write_text("stub")
monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
_prebuilt(monkeypatch, release_tag = "b9596", llama_tag = "b9596")
monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9596)
st = upd.get_update_status()
assert st["update_available"] is False
assert st["latest_tag"] == "b9596"
def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path):
# While the updater swaps the tree, status polls must not exec the binary
# being replaced (on Windows that exec can fail the installer's os.replace);
@ -420,8 +453,8 @@ def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path):
# --- installer-argument construction (mirrors the post-#5963 setup scripts) ---
def test_rocm_install_args_lemonade_gfx():
# Lemonade HIP app bundle: gfx family lives in the asset name.
def test_rocm_install_args_gfx_family():
# Per-gfx ROCm bundle: gfx family lives in the asset name.
assert upd._rocm_install_args("app-b9585-linux-x64-rocm-gfx110X.tar.gz") == [
"--rocm-gfx",
"gfx110x",

87
studio/backend/tests/test_llama_route_timeouts.py Normal file
View file

@ -0,0 +1,87 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import asyncio
import os
import sys
import time
from types import SimpleNamespace
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
import routes.inference as inf_mod # noqa: E402
def test_non_streaming_generation_timeout_has_read_deadline():
timeout = inf_mod._llama_non_streaming_generation_timeout()
assert timeout.read == inf_mod._DEFAULT_FIRST_TOKEN_TIMEOUT_S
def test_stream_first_item_deadline_after_headers():
async def _run():
class _Never:
async def __anext__(self):
await asyncio.Future()
started = time.monotonic()
try:
async for _ in inf_mod._aiter_llama_stream_items(
_Never(),
first_token_deadline = started + 0.02,
):
pass
except inf_mod.httpx.ReadTimeout:
pass
else:
raise AssertionError("first item deadline did not fire")
assert time.monotonic() - started < 0.5
asyncio.run(_run())
def test_preheader_send_cleanup_on_disconnect_and_cancel():
async def _run(cancel_parent):
state = SimpleNamespace(disconnected = False, closed = False, cancelled = False)
started = asyncio.Event()
class _Client:
async def send(
self,
req,
stream = False,
):
started.set()
try:
await asyncio.Future()
except asyncio.CancelledError:
state.cancelled = True
raise
async def aclose(self):
state.closed = True
class _Request:
async def is_disconnected(self):
return state.disconnected
task = asyncio.create_task(
inf_mod._send_stream_with_preheader_cancel(_Client(), object(), request = _Request())
)
await started.wait()
if cancel_parent:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
else:
raise AssertionError("helper cancellation did not propagate")
else:
state.disconnected = True
assert await task is None
assert state.closed
assert state.cancelled
asyncio.run(_run(False))
asyncio.run(_run(True))

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

@ -25,8 +25,12 @@ _spec.loader.exec_module(_lsa)
is_managed_flag = _lsa.is_managed_flag
parse_cache_override = _lsa.parse_cache_override
parse_ctx_override = _lsa.parse_ctx_override
parse_split_mode_override = _lsa.parse_split_mode_override
resolve_cache_type_kv = _lsa.resolve_cache_type_kv
resolve_tensor_parallel = _lsa.resolve_tensor_parallel
strip_shadowing_flags = _lsa.strip_shadowing_flags
strip_split_mode_only = _lsa.strip_split_mode_only
extra_args_disable_mmproj = _lsa.extra_args_disable_mmproj
validate_extra_args = _lsa.validate_extra_args
@ -509,6 +513,128 @@ def test_strip_shadowing_flags_defaults_strip_everything():
assert out == []
# ── --split-mode (Tensor Parallelism toggle) ─────────────────────────
# Soft-shadowed exactly like --cache-type-*: pass-through allowed (keeps
# the row/none/layer modes the boolean toggle doesn't expose), stripped
# on inherit, and reconciled back into the round-tripped tensor_parallel
# state.
@pytest.mark.parametrize(
"args",
[
["--split-mode", "tensor"],
["--split-mode", "row"],
["--split-mode", "none"],
["--split-mode", "layer"],
["-sm", "tensor"],
["--split-mode=row"],
["-sm=tensor"],
],
)
def test_split_mode_passes_through(args):
# Not denylisted -- a user keeps row/none/layer via extras.
assert validate_extra_args(args) == args
def test_split_mode_is_not_managed():
assert is_managed_flag("--split-mode") is False
assert is_managed_flag("-sm") is False
@pytest.mark.parametrize(
"args,expected",
[
(None, None),
([], None),
(["--top-k", "20"], None),
(["--split-mode", "tensor"], "tensor"),
(["--split-mode", "row"], "row"),
(["-sm", "none"], "none"),
(["--split-mode=layer"], "layer"),
(["-sm=tensor"], "tensor"),
# last-wins when supplied twice
(["-sm", "row", "--split-mode", "tensor"], "tensor"),
],
)
def test_parse_split_mode_override(args, expected):
assert parse_split_mode_override(args) == expected
@pytest.mark.parametrize(
"args",
[
["--split-mode"],
["-sm"],
["--split-mode", "-c", "4096"], # next token is a flag, not a value
],
)
def test_parse_split_mode_override_rejects_malformed_values(args):
with pytest.raises(ValueError, match = "split-mode|'-sm'"):
parse_split_mode_override(args)
def test_validate_extra_args_rejects_malformed_split_mode():
# Validation catches a value-less --split-mode at the boundary,
# mirroring the early --ctx-size / --cache-type checks.
with pytest.raises(ValueError, match = "split-mode"):
validate_extra_args(["--split-mode"])
@pytest.mark.parametrize(
"args,fallback,expected",
[
# No override -> fall back to the toggle value, both directions.
(["--top-k", "20"], True, True),
(["--top-k", "20"], False, False),
(None, True, True),
([], False, False),
# Explicit override wins: tensor -> on, anything else -> off,
# regardless of the toggle fallback.
(["--split-mode", "tensor"], False, True),
(["-sm", "tensor"], False, True),
(["--split-mode", "row"], True, False),
(["--split-mode", "none"], True, False),
(["--split-mode", "layer"], True, False),
(["--split-mode=tensor"], False, True),
# Case-insensitive on the mode string.
(["--split-mode", "TENSOR"], False, True),
# last-wins across multiple --split-mode flags.
(["-sm", "tensor", "--split-mode", "row"], True, False),
],
)
def test_resolve_tensor_parallel(args, fallback, expected):
assert resolve_tensor_parallel(args, fallback) is expected
def test_strip_shadowing_flags_drops_split_mode_when_requested():
out = strip_shadowing_flags(
["--split-mode", "row", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = True,
)
assert out == ["--top-k", "20"]
def test_extra_args_disable_mmproj_detects_flag():
assert extra_args_disable_mmproj(["--no-mmproj"]) is True
assert extra_args_disable_mmproj(["--threads", "12", "--no-mmproj"]) is True
assert extra_args_disable_mmproj(["--no-mmproj-auto"]) is True
def test_extra_args_disable_mmproj_false_when_absent():
assert extra_args_disable_mmproj(None) is False
assert extra_args_disable_mmproj(["--threads", "12"]) is False
def test_extra_args_disable_mmproj_last_wins():
assert extra_args_disable_mmproj(["--no-mmproj", "--mmproj-auto"]) is False
assert extra_args_disable_mmproj(["--mmproj-auto", "--no-mmproj-auto"]) is True
def test_strip_shadowing_flags_drops_model_draft_with_spec():
# --model-draft (and aliases) are Studio-managed since the separate
# MTP drafter support: an inherited copy must not last-wins-override
@ -523,6 +649,86 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec():
assert out == ["--top-k", "20"]
def test_strip_shadowing_flags_keeps_split_mode_when_not_requested():
# No tensor_parallel field supplied on the Apply -> an inherited
# --split-mode survives (mirrors the chat-template keep behavior).
out = strip_shadowing_flags(
["--split-mode", "row", "--top-k", "20"],
strip_context = True,
strip_cache = True,
strip_spec = True,
strip_template = True,
strip_split_mode = False,
)
assert out == ["--split-mode", "row", "--top-k", "20"]
def test_strip_shadowing_flags_drops_split_mode_short_alias_and_equals():
assert strip_shadowing_flags(["-sm", "tensor", "--top-k", "20"], strip_split_mode = True) == [
"--top-k",
"20",
]
assert strip_shadowing_flags(["--split-mode=row", "--seed", "-1"], strip_split_mode = True) == [
"--seed",
"-1",
]
def test_strip_shadowing_flags_defaults_strip_split_mode_too():
# The route's already-loaded comparator (no kwargs) must see a stored
# --split-mode as a shadowing flag so it forces a reload.
assert strip_shadowing_flags(["--split-mode", "tensor"]) == []
@pytest.mark.parametrize(
"args",
[
["--split-mode", "tensor", "-c", "4096"],
["-sm", "tensor", "-c", "4096"],
["--split-mode=tensor", "-c", "4096"],
["-sm=tensor", "-c", "4096"],
],
)
def test_strip_split_mode_only_keeps_other_shadow_flags(args):
# Every --split-mode form (long/short, space/=) is dropped; -c survives.
assert strip_split_mode_only(args) == ["-c", "4096"]
def test_strip_split_mode_only_preserves_none_and_empty():
# None means "inherit"; [] means "explicit empty" -- both must round-trip.
assert strip_split_mode_only(None) is None
assert strip_split_mode_only([]) == []
def test_strip_shadowing_flags_drops_tensor_split_with_split_mode():
# --tensor-split is coupled to the split mode: stripped together so a stale
# ratio can't override Studio's computed tensor split. Other flags survive.
out = strip_shadowing_flags(
["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = True,
)
assert out == ["--top-k", "20"]
def test_strip_shadowing_flags_keeps_tensor_split_when_not_requested():
# strip_split_mode=False keeps the whole split group (mode + ratios).
assert strip_shadowing_flags(
["--tensor-split", "1,1", "--top-k", "20"], strip_split_mode = False
) == ["--tensor-split", "1,1", "--top-k", "20"]
def test_strip_split_mode_only_drops_tensor_split_too():
# Downgrade / layer fallback must drop the coupled --tensor-split (all forms).
assert strip_split_mode_only(
["--split-mode", "tensor", "--tensor-split", "1,1", "-c", "4096"]
) == ["-c", "4096"]
assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == []
def test_strip_shadowing_flags_keeps_model_draft_without_spec():
out = strip_shadowing_flags(
["--model-draft", "/custom/mtp.gguf"],

611
studio/backend/tests/test_mcp_servers.py
View file

@ -598,3 +598,614 @@ def test_safetensors_agentic_empty_allowlist_still_means_allow_all():
)
# Empty allow-list = run anything (preserved contract).
assert calls == [("python", {"code": "1"})] or len(calls) >= 1
# ── discovery cache ─────────────────────────────────────────────────
def _one_tool(name = "echo"):
return [{"name": name, "inputSchema": {"type": "object", "properties": {}}}]
def test_get_enabled_mcp_tools_caches_discovery(tmp_path, monkeypatch):
"""A second send must serve tools from cache instead of re-probing."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
calls: list[str] = []
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
calls.append(url)
return _one_tool()
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
first = asyncio.run(tools_mod.get_enabled_mcp_tools())
second = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert len(calls) == 1 # probed once, cache hit on the second send
assert [t["function"]["name"] for t in first] == ["mcp__s1__echo"]
assert first == second
def test_get_enabled_mcp_tools_does_not_cache_failures(tmp_path, monkeypatch):
"""A failed probe isn't cached: once the cool-off elapses, it's retried."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
attempts = {"n": 0}
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
attempts["n"] += 1
if attempts["n"] == 1:
raise RuntimeError("server down")
return _one_tool()
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # failure -> empty
# Expire the cool-off (an until-time in the past) so the server is retried.
mcp_client._probe_cooloff_until["s1"] = 0.0
second = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert attempts["n"] == 2 # retried after the cool-off, not cached
assert [t["function"]["name"] for t in second] == ["mcp__s1__echo"]
def test_refresh_warms_tool_cache(tmp_path, monkeypatch):
"""Clicking Refresh must populate the cache the chat path reads."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
async def fake_refresh(
url,
headers = None,
timeout = None,
use_oauth = False,
):
return _one_tool()
monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh)
res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u"))
assert res.ok and res.tool_count == 1
def boom(*a, **k):
raise AssertionError("chat path re-probed despite a warm cache")
monkeypatch.setattr(tools_mod, "list_tools_async", boom)
specs = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert [t["function"]["name"] for t in specs] == ["mcp__s1__echo"]
def test_update_url_evicts_tool_cache(tmp_path, monkeypatch):
"""Re-pointing the URL must drop the old endpoint's cached tools."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("stale")})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True)
asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "https://new/mcp"), current_subject = "u"
)
)
assert mcp_client.get_cached_tools("s1") is None
def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch):
"""A rename touches no endpoint, so the cache must survive it."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
cached = _one_tool()
monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": cached})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
asyncio.run(
routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u")
)
assert mcp_client.get_cached_tools("s1") == cached
def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch):
"""Disabling a server must drop its cached tools, not leave them unread."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
asyncio.run(
routes_mcp.update_mcp_server("s1", McpServerUpdate(is_enabled = False), current_subject = "u")
)
assert mcp_client.get_cached_tools("s1") is None
def test_delete_evicts_tool_cache(tmp_path, monkeypatch):
"""Deleting a server must not leave its tools cached."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
asyncio.run(routes_mcp.delete_mcp_server("s1", current_subject = "u"))
assert mcp_client.get_cached_tools("s1") is None
def test_invalidate_tool_cache_clears_all(monkeypatch):
from core.inference import mcp_client
monkeypatch.setattr(mcp_client, "_tool_cache", {"a": _one_tool(), "b": _one_tool()})
mcp_client.invalidate_tool_cache()
assert mcp_client.get_cached_tools("a") is None
assert mcp_client.get_cached_tools("b") is None
def test_get_enabled_mcp_tools_probes_only_uncached(tmp_path, monkeypatch):
"""An already-cached server must not be re-probed alongside a cold one."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("cached")})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", is_enabled = True)
mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://b/mcp", is_enabled = True)
probed: list[str] = []
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
probed.append(url)
return _one_tool("fresh")
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
specs = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert probed == ["https://b/mcp"] # only the uncached server is probed
assert sorted(t["function"]["name"] for t in specs) == ["mcp__s1__cached", "mcp__s2__fresh"]
def test_get_enabled_mcp_tools_partial_failure_caches_healthy(tmp_path, monkeypatch):
"""One server failing must not stop the others from being cached/served."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://bad/mcp", is_enabled = True)
mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://good/mcp", is_enabled = True)
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
if "bad" in url:
raise RuntimeError("down")
return _one_tool("ok")
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
specs = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert [t["function"]["name"] for t in specs] == ["mcp__s2__ok"]
assert mcp_client.get_cached_tools("s1") is None # failure not cached
assert mcp_client.get_cached_tools("s2") == _one_tool("ok") # healthy cached
def test_get_enabled_mcp_tools_caches_empty_tool_list(tmp_path, monkeypatch):
"""A server exposing zero tools is cached as [] (a hit), not re-probed."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
calls: list[str] = []
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
calls.append(url)
return []
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == []
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == []
assert len(calls) == 1 # [] is a cache hit, not re-probed every send
assert mcp_client.get_cached_tools("s1") == []
def test_update_headers_evicts_tool_cache(tmp_path, monkeypatch):
"""Changing auth headers must drop tools discovered under the old headers."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(headers = {"Authorization": "Bearer new"}),
current_subject = "u",
)
)
assert mcp_client.get_cached_tools("s1") is None
def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe(tmp_path, monkeypatch):
"""A config edit landing during an in-flight probe must not be clobbered
by the now-stale probe result (TOCTOU on the cache write)."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True)
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
# Simulate a PUT landing while we are awaiting the probe.
mcp_servers_db.update_server("s1", {"url": "https://new/mcp"})
return _one_tool()
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
specs = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert specs == [] # stale result is neither served...
assert mcp_client.get_cached_tools("s1") is None # ...nor cached
def test_get_enabled_mcp_tools_no_cooloff_when_config_changes_mid_failed_probe(
tmp_path, monkeypatch
):
"""An edit landing while a probe of the OLD config is failing must not park
a cool-off on the now-fresh config -- else the re-pointed server the user
just fixed is needlessly skipped for the whole cool-off window."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True)
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
# The user re-points the server while the old endpoint's probe fails.
mcp_servers_db.update_server("s1", {"url": "https://new/mcp"})
raise RuntimeError("old endpoint down")
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == []
# The failure was for the OLD config, so the new one must stay re-probable.
assert not mcp_client.in_failure_cooloff("s1")
def test_get_enabled_mcp_tools_no_cooloff_when_server_deleted_mid_failed_probe(
tmp_path, monkeypatch
):
"""A delete landing while a probe fails must not leave an orphan cool-off
entry keyed by the since-removed server id."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
mcp_servers_db.delete_server("s1")
raise RuntimeError("down")
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == []
assert "s1" not in mcp_client._probe_cooloff_until # no orphan cool-off
def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff(tmp_path, monkeypatch):
"""A down server is probed once, then skipped during the cool-off instead
of being re-probed (and re-hung) on every send."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
attempts = {"n": 0}
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
attempts["n"] += 1
raise RuntimeError("down")
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # probes, fails
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # within cool-off
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # still skipped
assert attempts["n"] == 1 # only the first send probed
def test_cache_tools_clears_failure_cooloff(monkeypatch):
"""A successful probe lifts a server's failure cool-off."""
from core.inference import mcp_client
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_client.record_probe_failure("s1")
assert mcp_client.in_failure_cooloff("s1")
mcp_client.cache_tools("s1", _one_tool())
assert not mcp_client.in_failure_cooloff("s1")
def test_oauth_failure_cools_off_longer_than_plain(monkeypatch):
"""An OAuth server's failure cools off longer than a plain server's, so its
multi-minute probe hang doesn't recur every minute."""
from core.inference import mcp_client
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_client.record_probe_failure("plain", use_oauth = False)
mcp_client.record_probe_failure("oauth", use_oauth = True)
assert mcp_client._probe_cooloff_until["oauth"] > mcp_client._probe_cooloff_until["plain"]
def test_invalidate_clears_failure_cooloff(monkeypatch):
"""Eviction drops the failure cool-off so an edited server re-probes at once."""
from core.inference import mcp_client
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {"s1": 1.0, "s2": 2.0})
mcp_client.invalidate_tool_cache("s1")
assert "s1" not in mcp_client._probe_cooloff_until
assert "s2" in mcp_client._probe_cooloff_until
mcp_client.invalidate_tool_cache()
assert mcp_client._probe_cooloff_until == {}
def test_refresh_failure_records_cooloff(tmp_path, monkeypatch):
"""A failed manual refresh starts the cool-off so the next chat send does
not immediately hang on the down server."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
async def boom(
url,
headers = None,
timeout = None,
use_oauth = False,
):
raise RuntimeError("down")
monkeypatch.setattr(routes_mcp, "list_tools_async", boom)
res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u"))
assert res.ok is False
assert mcp_client.in_failure_cooloff("s1")
def test_refresh_drops_result_when_config_changes_mid_probe(tmp_path, monkeypatch):
"""A manual refresh must not warm the chat cache with tools discovered
under an old config if the server is edited while the probe is in flight."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True)
async def fake_refresh(
url,
headers = None,
timeout = None,
use_oauth = False,
):
mcp_servers_db.update_server("s1", {"url": "https://new/mcp"})
return _one_tool("stale")
monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh)
res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u"))
assert res.ok and res.tool_count == 1
assert mcp_client.get_cached_tools("s1") is None
def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monkeypatch):
"""A manual refresh failure for an old config must not cool off the freshly
edited server."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True)
async def boom(
url,
headers = None,
timeout = None,
use_oauth = False,
):
mcp_servers_db.update_server("s1", {"url": "https://new/mcp"})
raise RuntimeError("old endpoint down")
monkeypatch.setattr(routes_mcp, "list_tools_async", boom)
res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u"))
assert res.ok is False
assert not mcp_client.in_failure_cooloff("s1")
def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe(tmp_path, monkeypatch):
"""A delete landing while a probe is in flight must drop the now-orphan
result -- the `fresh is None` arm of the mid-probe TOCTOU guard. The
result is neither served nor cached under the since-removed id."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True)
async def fake(
url,
headers = None,
timeout = None,
use_oauth = False,
):
# Simulate a DELETE landing while we await the probe.
mcp_servers_db.delete_server("s1")
return _one_tool()
monkeypatch.setattr(tools_mod, "list_tools_async", fake)
specs = asyncio.run(tools_mod.get_enabled_mcp_tools())
assert specs == [] # orphan result not served
assert mcp_client.get_cached_tools("s1") is None # nor cached under a gone id
def test_oauth_probe_failure_in_chat_path_uses_long_cooloff(tmp_path, monkeypatch):
"""When an OAuth server fails discovery during a send, the chat path must
record the OAuth (long) cool-off, not the plain one -- otherwise its
multi-minute browser hang recurs every minute."""
import asyncio
import time
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from core.inference import tools as tools_mod
monkeypatch.setattr(mcp_client, "_tool_cache", {})
monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {})
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "https://x/mcp",
is_enabled = True,
use_oauth = True,
)
async def boom(
url,
headers = None,
timeout = None,
use_oauth = False,
):
raise RuntimeError("oauth down")
monkeypatch.setattr(tools_mod, "list_tools_async", boom)
assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == []
assert mcp_client.in_failure_cooloff("s1")
# The recorded window must exceed the plain cool-off, proving the OAuth
# branch (use_oauth=True) fired -- not the 60 s default.
remaining = mcp_client._probe_cooloff_until["s1"] - time.monotonic()
assert remaining > mcp_client.FAILED_PROBE_COOLOFF_SECONDS

164
studio/backend/tests/test_mcp_stdio_pr5863.py
View file

@ -7,6 +7,7 @@ and reaches it when enabled. The transport is stubbed so no subprocess spawns;
a recorder asserts whether it was reached.
"""
import os
import sys
import pytest
@ -14,11 +15,16 @@ from fastapi import HTTPException
from core.inference import mcp_client
from storage import mcp_servers_db
from utils import host_policy
def _reset_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
# The discovered-tool cache is process-global and keyed by server id; tests
# reuse "stdio1", so clear it (and the failure cool-off) for isolation —
# otherwise a prior test's warm cache makes discovery skip its probe.
mcp_client.invalidate_tool_cache()
def _enable(monkeypatch):
@ -29,6 +35,25 @@ def _disable(monkeypatch):
monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
@pytest.fixture(autouse = True)
def _isolate_stdio_env():
# apply_stdio_mcp_loopback_default() mutates os.environ and a module flag that
# monkeypatch can't roll back, and stdio_mcp_enabled() reads the process tool
# policy; snapshot/restore all three so nothing leaks between tests or files.
from state import tool_policy
saved = os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP")
saved_policy = tool_policy.get_tool_policy()
host_policy._reset_loopback_default_state()
yield
host_policy._reset_loopback_default_state()
tool_policy.set_tool_policy(saved_policy)
if saved is None:
os.environ.pop("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", None)
else:
os.environ["UNSLOTH_STUDIO_ALLOW_STDIO_MCP"] = saved
# ── transport stub + recorder ───────────────────────────────────────
@ -177,6 +202,145 @@ def test_stdio_enabled_only_for_exact_one(monkeypatch):
assert mcp_client.stdio_mcp_enabled() is True
# ── 3b. loopback bind defaults the gate on ──────────────────────────
@pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "LOCALHOST", "::1"])
def test_is_external_host_false_for_loopback(host):
assert host_policy.is_external_host(host) is False
# 127.0.0.2 is loopback in principle, but the rest of the stack hard-codes
# 127.0.0.1, so only the exact aliases count as local here.
@pytest.mark.parametrize("host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.10", "example.com"])
def test_is_external_host_true_for_network(host):
assert host_policy.is_external_host(host) is True
@pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "LOCALHOST", "::1"])
def test_loopback_bind_enables_stdio(monkeypatch, host):
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default(host)
assert mcp_client.stdio_mcp_enabled() is True
@pytest.mark.parametrize("host", ["0.0.0.0", "::", "127.0.0.2", "192.168.1.10", "example.com"])
def test_network_bind_leaves_stdio_off(monkeypatch, host):
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default(host)
assert mcp_client.stdio_mcp_enabled() is False
def test_colab_loopback_does_not_auto_enable(monkeypatch):
# Colab loopback is a hosted VM reachable via the proxy, so it stays off.
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1", is_colab = True)
assert mcp_client.stdio_mcp_enabled() is False
def test_explicit_enable_survives_colab(monkeypatch):
# An explicit operator opt-in still wins over the Colab exclusion (apply_
# early-returns on an explicit value, before the is_colab check).
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1", is_colab = True)
assert mcp_client.stdio_mcp_enabled() is True
def test_explicit_disable_survives_loopback(monkeypatch):
# An explicit =0 must not be overridden by the loopback auto-default.
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "0")
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1")
assert mcp_client.stdio_mcp_enabled() is False
def test_explicit_enable_survives_network_bind(monkeypatch):
# A deliberate network opt-in (-H 0.0.0.0 + var=1) must not be clobbered.
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
host_policy.apply_stdio_mcp_loopback_default("0.0.0.0")
assert mcp_client.stdio_mcp_enabled() is True
def test_loopback_default_not_inherited_by_later_public_bind(monkeypatch):
# Reusing run_server in one process: a loopback launch auto-enables, a later
# 0.0.0.0 launch must take it back down (not inherit it as an opt-in).
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1")
assert mcp_client.stdio_mcp_enabled() is True
host_policy.apply_stdio_mcp_loopback_default("0.0.0.0")
assert mcp_client.stdio_mcp_enabled() is False
@pytest.mark.parametrize("second_host", ["127.0.0.1", "0.0.0.0"])
def test_force_disable_after_auto_default_in_same_process(monkeypatch, second_host):
# Reuse: a loopback launch auto-enables, then the operator sets =0 before a
# later launch. The force-disable must win whether the later bind is loopback
# (must not rewrite to 1) or public (the relinquish path must not pop the =0).
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1")
assert mcp_client.stdio_mcp_enabled() is True
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "0")
host_policy.apply_stdio_mcp_loopback_default(second_host)
assert mcp_client.stdio_mcp_enabled() is False
def test_cleared_env_after_auto_default_falls_back_to_host_default(monkeypatch):
# Unsetting the var (unlike =0) is "no preference", so a loopback re-apply
# re-enables -- the asymmetry the staleness guard documents.
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1")
monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1")
assert mcp_client.stdio_mcp_enabled() is True
def test_disable_tools_overrides_loopback_default(monkeypatch):
# When stdio is on only via the loopback auto-default, --disable-tools (the
# only way tool policy is False on a loopback bind) turns it back off.
from state import tool_policy
_disable(monkeypatch)
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1")
assert mcp_client.stdio_mcp_enabled() is True
tool_policy.set_tool_policy(False)
assert mcp_client.stdio_mcp_enabled() is False
def test_explicit_env_opt_in_survives_external_default_policy(monkeypatch):
# `UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 unsloth studio run -H 0.0.0.0` with no
# --enable-tools: tool policy is False by the external-host default, not by
# --disable-tools, so the explicit env opt-in must still win.
from state import tool_policy
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
host_policy.apply_stdio_mcp_loopback_default("0.0.0.0") # no-op: value is explicit
tool_policy.set_tool_policy(False)
assert mcp_client.stdio_mcp_enabled() is True
def test_explicit_env_opt_in_beats_disable_tools_on_loopback(monkeypatch):
# An operator who hand-sets =1 before launch outranks --disable-tools even on
# loopback: apply_ leaves the auto-default inactive, so the veto doesn't apply.
from state import tool_policy
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
host_policy.apply_stdio_mcp_loopback_default("127.0.0.1") # no-op: value is explicit
tool_policy.set_tool_policy(False)
assert mcp_client.stdio_mcp_enabled() is True
@pytest.mark.parametrize("policy", [None, True])
def test_non_false_tool_policy_defers_to_env(monkeypatch, policy):
# Only an explicit --disable-tools (False) gates stdio; None/True fall through
# to the env var so the gate keeps its normal meaning.
from state import tool_policy
tool_policy.set_tool_policy(policy)
_disable(monkeypatch)
assert mcp_client.stdio_mcp_enabled() is False
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
assert mcp_client.stdio_mcp_enabled() is True
# ── 4. probe_timeout ────────────────────────────────────────────────

7
studio/backend/tests/test_mlx_training_worker_config.py
View file

@ -85,6 +85,13 @@ def test_mlx_studio_rejects_unknown_scheduler():
_normalize_mlx_studio_scheduler("linear_typo")
def test_mlx_studio_keeps_hf_style_tokenizer_dual_purpose():
source = (Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py").read_text()
assert "tokenizer = tokenizer" in source
assert "processor = tokenizer if is_vlm else None" not in source
def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer():
assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256)
assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512)

40
studio/backend/tests/test_mmproj_vram_accounting.py Normal file
View file

@ -0,0 +1,40 @@
# 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 mmproj VRAM accounting in GGUF fit budgeting (#5825)."""
from __future__ import annotations
from pathlib import Path
from core.inference.llama_cpp import LlamaCppBackend
def _write(path: Path, n_bytes: int) -> Path:
path.parent.mkdir(parents = True, exist_ok = True)
path.write_bytes(b"\x00" * n_bytes)
return path
def _backend() -> LlamaCppBackend:
return LlamaCppBackend.__new__(LlamaCppBackend)
def test_counts_resolved_projector_size(tmp_path: Path):
mmproj = _write(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf", 1024)
got = _backend()._mmproj_vram_bytes(str(mmproj))
assert got == 1024
def test_zero_when_no_projector_resolved(tmp_path: Path):
assert _backend()._mmproj_vram_bytes(None) == 0
def test_zero_when_projector_missing_on_disk(tmp_path: Path):
missing = tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf" # never created
got = _backend()._mmproj_vram_bytes(str(missing))
assert got == 0

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

@ -1,13 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through.
Covers ChatMessage tool/assistant roles, ChatCompletionRequest tool fields and
extra="allow", anthropic_tool_choice_to_openai, _build_passthrough_payload
tool_choice propagation, and _friendly_error's httpx-to-"Lost connection"
mapping. No server or GPU required.
"""
"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through."""
import os
import sys
@ -28,11 +22,13 @@ from models.inference import (
ChatMessage,
CompletionChoice,
CompletionMessage,
ResponsesRequest,
)
from core.inference.anthropic_compat import (
anthropic_tool_choice_to_openai,
)
from routes.inference import (
_build_chat_request,
_build_openai_passthrough_body,
_build_passthrough_payload,
_clamp_finish_reason,
@ -401,6 +397,28 @@ class TestChatCompletionRequestToolFields:
)
self._assert_unsupported_n(resp)
def test_confirm_tool_calls_rejected_for_provider_tools(self, monkeypatch):
class _UnusedBackend:
is_loaded = False
client = self._v1_client(monkeypatch, _UnusedBackend())
resp = client.post(
"/v1/chat/completions",
json = {
"messages": [{"role": "user", "content": "hi"}],
"provider_type": "openai",
"external_model": "gpt-4.1",
"enable_tools": True,
"enabled_tools": ["web_search"],
"confirm_tool_calls": True,
},
)
assert resp.status_code == 400
body = resp.json()
assert body["error"]["param"] == "confirm_tool_calls"
assert "only supported for local streaming tools" in body["error"]["message"]
def test_logprobs_rejected_until_supported(self, monkeypatch):
class _UnusedBackend:
is_loaded = False
@ -480,6 +498,7 @@ class TestChatCompletionRequestToolFields:
def test_n_rejected_for_non_gguf_path(self, monkeypatch):
class _NoGGUFBackend:
is_loaded = False
supports_tools = False
class _InferenceBackend:
active_model_name = "test-model"
@ -495,6 +514,45 @@ class TestChatCompletionRequestToolFields:
)
self._assert_unsupported_n(resp)
def test_confirm_tool_calls_requires_streaming_for_safetensors_tools(self, monkeypatch):
import routes.inference as inference_route
class _NoGGUFBackend:
is_loaded = False
supports_tools = False
class _InferenceBackend:
active_model_name = "test-model"
models = {"test-model": {"chat_template_info": {"template": "chatml"}}}
def generate_chat_completion_with_tools(self, **kwargs):
raise AssertionError("tool loop should be rejected before starting")
def generate_chat_completion(self, **kwargs):
raise AssertionError("plain path should not be used")
monkeypatch.setattr(
inference_route,
"_detect_safetensors_features",
lambda backend, chat_template: {"supports_tools": True},
)
client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend())
resp = client.post(
"/v1/chat/completions",
json = {
"messages": [{"role": "user", "content": "hi"}],
"enable_tools": True,
"enabled_tools": ["web_search"],
"confirm_tool_calls": True,
"stream": False,
},
)
assert resp.status_code == 400
body = resp.json()
assert body["error"]["param"] == "confirm_tool_calls"
assert "requires stream=true" in body["error"]["message"]
def test_multiturn_tool_loop_messages(self):
req = ChatCompletionRequest(
messages = [
@ -724,6 +782,22 @@ class TestPassthroughReasoningKwargs:
)
assert body["chat_template_kwargs"] == {"reasoning_effort": "high"}
def test_reasoning_effort_none_forwarded_for_effort_style_models(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = False, reasoning_effort = "none"),
backend_ctx = 4096,
llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"),
)
assert body["chat_template_kwargs"] == {"reasoning_effort": "none"}
def test_reasoning_effort_minimal_maps_to_low_for_effort_style_models(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = True, reasoning_effort = "minimal"),
backend_ctx = 4096,
llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"),
)
assert body["chat_template_kwargs"] == {"reasoning_effort": "low"}
def test_enable_thinking_maps_to_effort_for_effort_style_models(self):
body = _build_openai_passthrough_body(
self._payload(enable_thinking = False),
@ -834,12 +908,6 @@ class TestOpenAICompatibilityHelpers:
class TestFriendlyErrorHttpx:
"""When llama-server is down, httpx RequestError strings lack the
"Lost connection to llama-server" substring the sync path keys off, so the
old substring-only `_friendly_error` returned a useless generic message.
These tests pin the new isinstance-based mapping.
"""
def _req(self):
return httpx.Request("POST", "http://127.0.0.1:65535/v1/chat/completions")
@ -857,7 +925,7 @@ class TestFriendlyErrorHttpx:
def test_read_timeout_mapped(self):
exc = httpx.ReadTimeout("timed out", request = self._req())
assert "Lost connection" in _friendly_error(exc)
assert "first token within 20 minutes" in _friendly_error(exc)
def test_non_httpx_unchanged(self):
# Non-httpx exceptions still fall through to the substring heuristics
@ -1206,6 +1274,45 @@ class TestGgufVisionToolRouting:
assert captured["kwargs"]["disable_parallel_tool_use"] is True
def test_confirm_tool_calls_requires_streaming_for_gguf_tools(self, monkeypatch):
import routes.inference as inf_mod
def _plain(**kwargs):
raise AssertionError("plain GGUF path should not be used")
def _tools(**kwargs):
raise AssertionError("tool loop should be rejected before starting")
backend = SimpleNamespace(
is_loaded = True,
is_vision = False,
supports_tools = True,
model_identifier = "test-gguf",
generate_chat_completion = _plain,
generate_chat_completion_with_tools = _tools,
)
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
payload = ChatCompletionRequest(
model = "default",
enable_tools = True,
enabled_tools = ["web_search"],
confirm_tool_calls = True,
stream = False,
messages = [{"role": "user", "content": "search once"}],
)
with pytest.raises(HTTPException) as exc:
self._drive(
openai_chat_completions(
payload,
request = self._Request(),
current_subject = "test",
)
)
assert exc.value.status_code == 400
assert "requires stream=true" in exc.value.detail["error"]["message"]
def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch):
import routes.inference as inf_mod
@ -1295,3 +1402,47 @@ class TestGgufVisionToolRouting:
assert seen_seeds == expected
assert [choice["index"] for choice in body["choices"]] == [0, 1, 2]
# =====================================================================
# Responses API -> Chat Completions translation: chat_template_kwargs
# (e.g. {"enable_thinking": true}) sent via the Responses extra-body must
# reach the built ChatCompletionRequest's typed ``enable_thinking`` field,
# otherwise /v1/responses silently ignores reasoning control (issue #6198).
# =====================================================================
class TestResponsesChatTemplateKwargs:
_messages = [ChatMessage(role = "user", content = "What is 100 - 67?")]
def test_enable_thinking_lifted_from_extra_body(self):
payload = ResponsesRequest(
model = "qwen-local",
input = "What is 100 - 67?",
chat_template_kwargs = {"enable_thinking": True},
)
chat_req = _build_chat_request(payload, self._messages, stream = False)
assert chat_req.enable_thinking is True
def test_enable_thinking_false_lifted_from_extra_body(self):
payload = ResponsesRequest(
model = "qwen-local",
input = "hi",
chat_template_kwargs = {"enable_thinking": False},
)
chat_req = _build_chat_request(payload, self._messages, stream = True)
assert chat_req.enable_thinking is False
def test_no_chat_template_kwargs_leaves_enable_thinking_unset(self):
payload = ResponsesRequest(model = "qwen-local", input = "hi")
chat_req = _build_chat_request(payload, self._messages, stream = False)
assert chat_req.enable_thinking is None
def test_chat_template_kwargs_without_enable_thinking_is_ignored(self):
payload = ResponsesRequest(
model = "qwen-local",
input = "hi",
chat_template_kwargs = {"some_other_flag": True},
)
chat_req = _build_chat_request(payload, self._messages, stream = False)
assert chat_req.enable_thinking is None

129
studio/backend/tests/test_rag_ingestion.py
View file

@ -39,7 +39,7 @@ def test_ingestion_lifecycle_pending_to_completed(rag_home, stub_embeddings, tmp
conn = rag_db.get_connection()
try:
assert store.get_document(conn, doc_id)["status"] == "pending"
assert store.get_document(conn, doc_id)["status"] in {"pending", "running", "completed"}
finally:
conn.close()
@ -83,6 +83,133 @@ def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path):
conn.close()
def test_ingestion_dedupe_removes_duplicate_upload(rag_home, stub_embeddings):
from utils.paths import ensure_dir, rag_uploads_root
uploads = ensure_dir(rag_uploads_root())
first_path = uploads / "doc.txt"
duplicate_path = uploads / "copy.txt"
first_path.write_text("alpha bravo charlie", encoding = "utf-8")
duplicate_path.write_text("alpha bravo charlie", encoding = "utf-8")
scope = store.project_scope("P1")
doc_id, job_id = ingestion.start_ingestion(
scope,
None,
None,
"doc.txt",
str(first_path),
project_id = "P1",
)
_drain(job_id)
_wait_completed(job_id)
doc_id2, job_id2 = ingestion.start_ingestion(
scope,
None,
None,
"copy.txt",
str(duplicate_path),
project_id = "P1",
)
events = _drain(job_id2)
assert doc_id2 == doc_id
assert any(e.get("deduped") for e in events)
assert first_path.exists()
assert not duplicate_path.exists()
def test_ingestion_retry_replaces_failed_hash(rag_home, stub_embeddings):
from utils.paths import ensure_dir, rag_uploads_root
uploads = ensure_dir(rag_uploads_root())
old_path = uploads / "failed.txt"
retry_path = uploads / "retry.txt"
old_path.write_text("alpha bravo charlie", encoding = "utf-8")
retry_path.write_text("alpha bravo charlie", encoding = "utf-8")
scope = store.project_scope("P1")
sha = ingestion._sha256_file(str(old_path))
conn = rag_db.get_connection()
try:
failed_id = store.create_document(
conn,
scope = scope,
filename = "failed.txt",
sha256 = sha,
project_id = "P1",
status = "failed",
stored_path = str(old_path),
)
finally:
conn.close()
doc_id, job_id = ingestion.start_ingestion(
scope,
None,
None,
"retry.txt",
str(retry_path),
project_id = "P1",
)
events = _drain(job_id)
assert doc_id != failed_id
assert not any(e.get("deduped") for e in events)
assert not old_path.exists()
assert retry_path.exists()
status = _wait_completed(job_id)
assert status["status"] == "completed"
conn = rag_db.get_connection()
try:
assert store.get_document(conn, failed_id) is None
assert store.get_document(conn, doc_id)["status"] == "completed"
finally:
conn.close()
def test_delete_document_route_removes_stored_upload(rag_home):
from fastapi import FastAPI
from fastapi.testclient import TestClient
from auth.authentication import get_current_subject
from routes.rag import router
from utils.paths import ensure_dir, rag_uploads_root
upload = ensure_dir(rag_uploads_root()) / "delete-me.txt"
upload.write_text("alpha bravo", encoding = "utf-8")
scope = store.project_scope("P1")
conn = rag_db.get_connection()
try:
doc_id = store.create_document(
conn,
scope = scope,
filename = "delete-me.txt",
sha256 = "delete-route-sha",
project_id = "P1",
status = "completed",
stored_path = str(upload),
)
finally:
conn.close()
app = FastAPI()
app.include_router(router, prefix = "/api/rag")
app.dependency_overrides[get_current_subject] = lambda: "tester"
client = TestClient(app)
res = client.delete(f"/api/rag/documents/{doc_id}")
assert res.status_code == 200
assert not upload.exists()
conn = rag_db.get_connection()
try:
assert store.get_document(conn, doc_id) is None
finally:
conn.close()
def test_ingestion_delete_removes_all_rows(rag_home, stub_embeddings, tmp_path):
path = _write(tmp_path, "doc.txt", "alpha bravo charlie delta")
scope = store.kb_scope("K1")

593
studio/backend/tests/test_responses_tool_passthrough.py
View file

@ -36,6 +36,8 @@ import json
import httpx
import pytest
from fastapi import HTTPException
from fastapi.responses import JSONResponse
from pydantic import ValidationError
from models.inference import (
@ -46,6 +48,7 @@ from models.inference import (
ResponsesInputMessage,
ResponsesOutputFunctionCall,
ResponsesOutputMessage,
ResponsesOutputReasoning,
ResponsesOutputTextContent,
ResponsesOutputTextPart,
ResponsesRequest,
@ -58,7 +61,8 @@ from routes.inference import (
_build_chat_request,
_chat_tool_calls_to_responses_output,
_normalise_responses_input,
_responses_tool_output_text,
_responses_tool_output_content,
_responses_non_streaming,
_responses_stream,
_translate_responses_tool_choice_to_chat,
_translate_responses_tools_to_chat,
@ -284,6 +288,59 @@ class TestBuildChatRequest:
assert chat_req.parallel_tool_calls is False
def test_chat_template_kwargs_enable_thinking_true_is_lifted(self):
payload = ResponsesRequest(
input = "hi",
chat_template_kwargs = {"enable_thinking": True},
)
messages = [ChatMessage(role = "user", content = "hi")]
chat_req = _build_chat_request(payload, messages, stream = False)
assert chat_req.enable_thinking is True
def test_chat_template_kwargs_enable_thinking_false_is_lifted(self):
payload = ResponsesRequest(
input = "hi",
chat_template_kwargs = {"enable_thinking": False},
)
messages = [ChatMessage(role = "user", content = "hi")]
chat_req = _build_chat_request(payload, messages, stream = False)
assert chat_req.enable_thinking is False
def test_reasoning_effort_high_enables_local_thinking(self):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"})
messages = [ChatMessage(role = "user", content = "hi")]
chat_req = _build_chat_request(payload, messages, stream = False)
assert chat_req.reasoning_effort == "high"
assert chat_req.enable_thinking is True
def test_reasoning_effort_none_disables_local_thinking(self):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"})
messages = [ChatMessage(role = "user", content = "hi")]
chat_req = _build_chat_request(payload, messages, stream = False)
assert chat_req.reasoning_effort == "none"
assert chat_req.enable_thinking is False
def test_explicit_enable_thinking_false_disables_reasoning_effort(self):
payload = ResponsesRequest(
input = "hi",
reasoning = {"effort": "high"},
chat_template_kwargs = {"enable_thinking": False},
)
messages = [ChatMessage(role = "user", content = "hi")]
chat_req = _build_chat_request(payload, messages, stream = False)
assert chat_req.reasoning_effort == "none"
assert chat_req.enable_thinking is False
# =====================================================================
# _normalise_responses_input — multi-turn tool mapping
@ -379,20 +436,144 @@ class TestNormaliseResponsesInputWithTools:
assert sum(1 for m in msgs if m.role == "system") == 1
assert "A" in msgs[0].content and "B" in msgs[0].content
def test_content_array_output_serialised_to_json_string(self):
def test_content_array_text_output_flattens_to_tool_text(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [{"type": "output_text", "text": "ok"}],
"output": [{"type": "input_text", "text": "ok"}],
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].role == "tool"
# Content is serialised so llama-server sees a string.
assert json.loads(msgs[0].content) == [{"type": "output_text", "text": "ok"}]
assert msgs[0].content == "ok"
def test_content_array_image_output_becomes_multimodal_tool_content(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [
{"type": "input_text", "text": "see image"},
{
"type": "input_image",
"image_url": "data:image/png;base64,AAA",
"detail": "high",
},
],
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].role == "tool"
assert msgs[0].tool_call_id == "call_1"
assert msgs[0].model_dump(exclude_none = True)["content"] == [
{"type": "text", "text": "see image"},
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,AAA",
"detail": "high",
},
},
]
chat_req = _build_chat_request(payload, msgs, stream = False)
assert chat_req.model_dump(exclude_none = True)["messages"][0]["content"] == [
{"type": "text", "text": "see image"},
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,AAA",
"detail": "high",
},
},
]
def test_content_array_image_output_allows_original_detail(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [
{
"type": "input_image",
"image_url": "https://example.com/screenshot.png",
"detail": "original",
},
],
}
],
)
msgs = _normalise_responses_input(payload)
assert msgs[0].model_dump(exclude_none = True)["content"] == [
{
"type": "image_url",
"image_url": {
"url": "https://example.com/screenshot.png",
"detail": "original",
},
},
]
def test_content_array_file_id_image_output_rejected_clearly(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [
{"type": "input_text", "text": "see image"},
{"type": "input_image", "file_id": "file_abc"},
],
}
],
)
with pytest.raises(HTTPException) as exc:
_normalise_responses_input(payload)
assert exc.value.status_code == 400
assert "file_id" in str(exc.value.detail)
def test_content_array_file_output_rejected_clearly(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [
{"type": "input_text", "text": "see file"},
{
"type": "input_file",
"file_data": "data:application/pdf;base64,AAA",
"filename": "report.pdf",
},
],
}
],
)
with pytest.raises(HTTPException) as exc:
_normalise_responses_input(payload)
assert exc.value.status_code == 400
assert "input_file" in str(exc.value.detail)
def test_content_array_malformed_image_output_rejected_clearly(self):
payload = ResponsesRequest(
input = [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [{"type": "input_image", "detail": "high"}],
}
],
)
with pytest.raises(HTTPException) as exc:
_normalise_responses_input(payload)
assert exc.value.status_code == 400
assert "image_url" in str(exc.value.detail)
def test_empty_function_call_output_gets_no_output_sentinel(self):
payload = ResponsesRequest(
@ -485,8 +666,8 @@ class TestNormaliseResponsesInputWithTools:
assert msgs[0].content == "(no output)"
def test_tool_output_serializer_preserves_non_empty_text(self):
assert _responses_tool_output_text("done") == "done"
assert _responses_tool_output_text(" done ") == " done "
assert _responses_tool_output_content("done") == "done"
assert _responses_tool_output_content(" done ") == " done "
# =====================================================================
@ -544,6 +725,119 @@ class TestChatToolCallsToResponsesOutput:
assert items[0]["arguments"] == ""
# =====================================================================
# Non-streaming Responses adapter
# =====================================================================
class TestResponsesNonStreamingAdapter:
class _Request:
pass
@staticmethod
def _run_with_message(
monkeypatch,
message,
payload = None,
llama_backend = None,
):
import routes.inference as inf_mod
async def fake_chat_completions(chat_req, request):
return JSONResponse(
content = {
"model": "test-model",
"choices": [{"message": message}],
"usage": {"prompt_tokens": 2, "completion_tokens": 3},
}
)
monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions)
if llama_backend is not None:
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: llama_backend)
payload = payload or ResponsesRequest(input = "hi")
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_non_streaming(
payload, messages, TestResponsesNonStreamingAdapter._Request()
)
return json.loads(response.body.decode())
return asyncio.run(run())
def test_think_block_becomes_reasoning_item_before_message(self, monkeypatch):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"})
body = self._run_with_message(
monkeypatch,
{"content": "<think>plan</think>33"},
payload = payload,
)
assert [item["type"] for item in body["output"]] == ["reasoning", "message"]
assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}]
assert body["output"][0]["summary"] == []
assert body["output"][1]["content"][0]["text"] == "33"
assert "<think>" not in body["output"][1]["content"][0]["text"]
assert "</think>" not in body["output"][1]["content"][0]["text"]
def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch):
body = self._run_with_message(monkeypatch, {"content": "show <think>x</think> tags"})
assert [item["type"] for item in body["output"]] == ["message"]
assert body["output"][0]["content"][0]["text"] == "show <think>x</think> tags"
def test_non_reasoning_gguf_keeps_literal_think_tags_visible(self, monkeypatch):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"})
body = self._run_with_message(
monkeypatch,
{"content": "show <think>x</think> tags"},
payload = payload,
llama_backend = SimpleNamespace(
is_loaded = True,
reasoning_always_on = False,
supports_reasoning = False,
),
)
assert [item["type"] for item in body["output"]] == ["message"]
assert body["output"][0]["content"][0]["text"] == "show <think>x</think> tags"
def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch):
body = self._run_with_message(
monkeypatch,
{
"content": "33",
"reasoning_content": [
{"type": "reasoning_text", "text": "plan"},
{"type": "reasoning_text", "text": " next"},
],
},
)
assert [item["type"] for item in body["output"]] == ["reasoning", "message"]
assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan next"}]
assert body["output"][1]["content"][0]["text"] == "33"
def test_plain_content_remains_message_only(self, monkeypatch):
body = self._run_with_message(monkeypatch, {"content": "33"})
assert [item["type"] for item in body["output"]] == ["message"]
assert body["output"][0]["content"][0]["text"] == "33"
def test_reasoning_only_is_also_visible_message_text(self, monkeypatch):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"})
body = self._run_with_message(
monkeypatch,
{"content": "<think>plan</think>"},
payload = payload,
)
assert [item["type"] for item in body["output"]] == ["reasoning", "message"]
assert body["output"][0]["content"][0]["text"] == "plan"
assert body["output"][1]["content"][0]["text"] == "plan"
# =====================================================================
# Streaming Responses adapter
# =====================================================================
@ -570,6 +864,262 @@ class TestResponsesStreamAdapter:
if line.startswith(prefix)
]
@staticmethod
def _install_stream_mock(
monkeypatch,
chunks,
*,
supports_reasoning = True,
reasoning_always_on = False,
):
import routes.inference as inf_mod
def handler(request: httpx.Request) -> httpx.Response:
content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks)
content += "data: [DONE]\n\n"
return httpx.Response(
200,
content = content.encode(),
headers = {"content-type": "text/event-stream"},
)
transport = httpx.MockTransport(handler)
real_async_client = httpx.AsyncClient
def _client(*args, **kwargs):
return real_async_client(
transport = transport,
timeout = kwargs.get("timeout", 600),
)
monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client)
monkeypatch.setattr(
inf_mod,
"get_llama_cpp_backend",
lambda: SimpleNamespace(
is_loaded = True,
is_vision = False,
context_length = 4096,
base_url = "http://llama.test",
supports_reasoning = supports_reasoning,
reasoning_always_on = reasoning_always_on,
_request_reasoning_kwargs = (
lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None
),
),
)
def test_split_think_markers_stream_as_reasoning_and_visible_text(self, monkeypatch):
chunks = [
{"choices": [{"delta": {"content": "<thi"}}]},
{"choices": [{"delta": {"content": "nk>pla"}}]},
{"choices": [{"delta": {"content": "n</th"}}]},
{"choices": [{"delta": {"content": "ink>33"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"})
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert "".join(event["delta"] for event in reasoning_deltas) == "plan"
assert "".join(event["delta"] for event in text_deltas) == "33"
completed = self._payloads(lines, "response.completed")[0]
assert [item["type"] for item in completed["response"]["output"]] == [
"reasoning",
"message",
]
assert completed["response"]["output"][0]["content"][0]["text"] == "plan"
assert completed["response"]["output"][1]["content"][0]["text"] == "33"
def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch):
chunks = [
{"choices": [{"delta": {"content": "show <thi"}}]},
{"choices": [{"delta": {"content": "nk>x</think> tags"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
payload = ResponsesRequest(input = "hi", stream = True)
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert reasoning_deltas == []
assert "".join(event["delta"] for event in text_deltas) == "show <think>x</think> tags"
completed = self._payloads(lines, "response.completed")[0]
assert [item["type"] for item in completed["response"]["output"]] == ["message"]
assert completed["response"]["output"][0]["content"][0]["text"] == (
"show <think>x</think> tags"
)
def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch):
chunks = [
{"choices": [{"delta": {"content": "show <thi"}}]},
{"choices": [{"delta": {"content": "nk>x</think> tags"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks, supports_reasoning = False)
payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"})
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert reasoning_deltas == []
assert "".join(event["delta"] for event in text_deltas) == "show <think>x</think> tags"
completed = self._payloads(lines, "response.completed")[0]
assert [item["type"] for item in completed["response"]["output"]] == ["message"]
assert completed["response"]["output"][0]["content"][0]["text"] == (
"show <think>x</think> tags"
)
def test_reasoning_only_streams_as_visible_message_text(self, monkeypatch):
chunks = [
{"choices": [{"delta": {"content": "<think>plan</think>"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"})
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert "".join(event["delta"] for event in reasoning_deltas) == "plan"
assert "".join(event["delta"] for event in text_deltas) == "plan"
completed = self._payloads(lines, "response.completed")[0]
assert [item["type"] for item in completed["response"]["output"]] == [
"reasoning",
"message",
]
assert completed["response"]["output"][0]["content"][0]["text"] == "plan"
assert completed["response"]["output"][1]["content"][0]["text"] == "plan"
def test_structured_reasoning_content_streams_as_reasoning(self, monkeypatch):
chunks = [
{"choices": [{"delta": {"reasoning_content": "plan"}}]},
{"choices": [{"delta": {"content": "33"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
payload = ResponsesRequest(input = "hi", stream = True)
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert "".join(event["delta"] for event in reasoning_deltas) == "plan"
assert "".join(event["delta"] for event in text_deltas) == "33"
completed = self._payloads(lines, "response.completed")[0]
assert completed["response"]["output"][0]["type"] == "reasoning"
assert completed["response"]["output"][1]["type"] == "message"
def test_structured_reasoning_content_parts_stream_as_reasoning(self, monkeypatch):
chunks = [
{
"choices": [
{
"delta": {
"reasoning_content": {
"content": [
{"type": "reasoning_text", "text": "plan"},
{"type": "reasoning_text", "text": " next"},
]
}
}
}
]
},
{"choices": [{"delta": {"content": "33"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
payload = ResponsesRequest(input = "hi", stream = True)
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert "".join(event["delta"] for event in reasoning_deltas) == "plan next"
assert "".join(event["delta"] for event in text_deltas) == "33"
assert "reasoning_text" not in "".join(event["delta"] for event in reasoning_deltas)
completed = self._payloads(lines, "response.completed")[0]
assert completed["response"]["output"][0]["content"][0]["text"] == "plan next"
assert completed["response"]["output"][1]["content"][0]["text"] == "33"
def test_tool_first_stream_closes_items_in_output_index_order(self, monkeypatch):
chunks = [
{
"choices": [
{
"delta": {
"tool_calls": [
{
"index": 0,
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"},
}
]
}
}
]
},
{"choices": [{"delta": {"content": "done"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
payload = ResponsesRequest(input = "hi", stream = True)
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(payload, messages, self._Request())
return await self._collect(response)
lines = asyncio.run(run())
done_events = self._payloads(lines, "response.output_item.done")
assert [event["output_index"] for event in done_events] == [0, 1]
assert [event["item"]["type"] for event in done_events] == ["function_call", "message"]
completed = self._payloads(lines, "response.completed")[0]
assert [item["type"] for item in completed["response"]["output"]] == [
"function_call",
"message",
]
def test_requests_usage_and_caps_parallel_tool_calls(self, monkeypatch):
import routes.inference as inf_mod
@ -678,6 +1228,15 @@ class TestResponsesStreamAdapter:
class TestResponsesOutputFunctionCall:
def test_reasoning_output_item_serialises_full_reasoning_content(self):
item = ResponsesOutputReasoning(content = [{"type": "reasoning_text", "text": "plan"}])
d = item.model_dump()
assert d["type"] == "reasoning"
assert d["id"].startswith("rs_")
assert d["status"] == "completed"
assert d["summary"] == []
assert d["content"] == [{"type": "reasoning_text", "text": "plan"}]
def test_direct_construction(self):
fc = ResponsesOutputFunctionCall(
call_id = "call_1",
@ -778,6 +1337,26 @@ class TestCodexStyleRequestShapes:
assert len(req.input) == 3
assert isinstance(req.input[1], ResponsesUnknownInputItem)
def test_emitted_reasoning_item_replay_is_dropped_for_local_chat(self):
payload = ResponsesRequest(
input = [
{"role": "user", "content": "Hi"},
{
"type": "reasoning",
"id": "rs_1",
"summary": [],
"content": [{"type": "reasoning_text", "text": "plan"}],
},
{"role": "assistant", "content": "33"},
{"role": "user", "content": "Continue"},
],
)
msgs = _normalise_responses_input(payload)
assert [m.role for m in msgs] == ["user", "assistant", "user"]
assert all("plan" not in (m.content or "") for m in msgs if isinstance(m.content, str))
def test_unknown_content_part_type_accepted(self):
"""Unknown content-part types (e.g. future input_audio) validate as
ResponsesUnknownContentPart so the request doesn't 422."""

274
studio/backend/tests/test_s3_dataset.py Normal file
View file

@ -0,0 +1,274 @@
# 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 S3 dataset loader (core.training.s3_dataset).
boto3 is optional and may be absent in CI, so the S3 client is mocked: a fake
client provides a paginator over a synthetic bucket listing and writes files on
download_file. No network or real AWS credentials are involved.
"""
import importlib.util
import os
from pathlib import Path
import pytest
# Load the modules under test directly by path. Importing them through their
# packages (core.training / models) would execute heavy package __init__ chains
# (structlog, torch, …) that aren't needed for these unit tests.
_BACKEND = Path(__file__).resolve().parents[1]
def _load(mod_name, rel_path):
spec = importlib.util.spec_from_file_location(mod_name, _BACKEND / rel_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
s3_dataset = _load("s3_dataset", "core/training/s3_dataset.py")
S3Config = _load("models_training_s3", "models/training.py").S3Config
class _FakePaginator:
def __init__(self, keys):
self._keys = keys
def paginate(self, **kwargs):
prefix = kwargs.get("Prefix")
contents = [{"Key": k} for k in self._keys if prefix is None or k.startswith(prefix)]
# Emit in two pages to exercise pagination handling.
mid = len(contents) // 2
yield {"Contents": contents[:mid]}
yield {"Contents": contents[mid:]}
class _FakeS3Client:
def __init__(self, keys):
self._keys = keys
self.downloaded = []
def get_paginator(self, name):
assert name == "list_objects_v2"
return _FakePaginator(self._keys)
def download_file(self, bucket, key, local_path, **kwargs):
self.downloaded.append((bucket, key, local_path))
callback = kwargs.get("Callback")
if callback is not None:
callback(1)
with open(local_path, "w", encoding = "utf-8") as f:
f.write(f"content-of:{key}")
@pytest.fixture
def fake_client(monkeypatch):
"""Force boto3_available True and stub the client builder."""
keys = [
"datasets/train.parquet",
"datasets/extra.parquet",
"datasets/notes.txt", # filtered out (unsupported)
"datasets/subdir/", # directory placeholder, skipped
"other/ignore.parquet", # filtered out by prefix
]
client = _FakeS3Client(keys)
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
return client
def _cfg(**overrides):
base = {
"bucket": "my-bucket",
"region": "us-east-1",
"prefix": "datasets/",
"access_key_id": "AKIA_TEST",
"secret_access_key": "secret",
"use_iam_role": False,
}
base.update(overrides)
return base
def test_downloads_only_supported_files_under_prefix(fake_client, tmp_path):
files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
names = sorted(os.path.basename(f) for f in files)
# txt is unsupported, the directory placeholder is skipped, and the
# "other/" key is excluded by the prefix filter.
assert names == ["extra.parquet", "train.parquet"]
for f in files:
assert os.path.exists(f)
def test_allows_json_and_jsonl_family(monkeypatch, tmp_path):
client = _FakeS3Client(["datasets/train.json", "datasets/extra.jsonl"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
assert sorted(os.path.basename(f) for f in files) == ["extra.jsonl", "train.json"]
def test_ignores_common_json_metadata_files(monkeypatch, tmp_path):
client = _FakeS3Client(
[
"datasets/train.parquet",
"datasets/schema.json",
"datasets/metadata.json",
"datasets/dataset_info.json",
]
)
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
assert [os.path.basename(f) for f in files] == ["train.parquet"]
def test_raises_when_prefix_contains_mixed_formats(monkeypatch, tmp_path):
client = _FakeS3Client(["datasets/train.parquet", "datasets/stray.csv"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
with pytest.raises(ValueError, match = "mixed dataset formats"):
s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
assert client.downloaded == []
def test_raises_when_no_supported_files(monkeypatch, tmp_path):
client = _FakeS3Client(["datasets/readme.txt"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
with pytest.raises(ValueError, match = "No supported dataset files"):
s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
def test_raises_when_boto3_missing(monkeypatch, tmp_path):
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: False)
with pytest.raises(RuntimeError, match = "requires boto3"):
s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
def test_basename_collisions_are_disambiguated(monkeypatch, tmp_path):
# Two keys share a basename under different sub-prefixes.
client = _FakeS3Client(["datasets/a/train.parquet", "datasets/b/train.parquet"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
assert len(files) == 2
assert len(set(files)) == 2 # no overwrite
def test_basename_collision_skips_existing_generated_suffix(monkeypatch, tmp_path):
client = _FakeS3Client(
[
"datasets/a/train.parquet",
"datasets/b/train_1.parquet",
"datasets/c/train.parquet",
]
)
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path))
assert [os.path.basename(f) for f in files] == [
"train.parquet",
"train_1.parquet",
"train_2.parquet",
]
assert len(set(files)) == 3
assert (tmp_path / "train_1.parquet").read_text(encoding = "utf-8") == (
"content-of:datasets/b/train_1.parquet"
)
assert (tmp_path / "train_2.parquet").read_text(encoding = "utf-8") == (
"content-of:datasets/c/train.parquet"
)
def test_download_handle_cleans_owned_temp_dir(monkeypatch, tmp_path):
target_dir = tmp_path / "owned-download"
client = _FakeS3Client(["datasets/train.parquet"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir))
download = s3_dataset.prepare_s3_dataset_download(_cfg())
assert target_dir.exists()
assert download.files == [str(target_dir / "train.parquet")]
download.cleanup()
assert not target_dir.exists()
def test_dest_dir_is_not_removed_by_cleanup(monkeypatch, tmp_path):
client = _FakeS3Client(["datasets/train.parquet"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
download = s3_dataset.prepare_s3_dataset_download(_cfg(), dest_dir = str(tmp_path))
download.cleanup()
assert tmp_path.exists()
assert (tmp_path / "train.parquet").exists()
def test_cancel_callback_aborts_and_removes_temp_dir(monkeypatch, tmp_path):
target_dir = tmp_path / "cancelled-download"
client = _FakeS3Client(["datasets/train.parquet"])
monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True)
monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client)
monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir))
calls = 0
def cancel_after_download_starts():
nonlocal calls
calls += 1
return calls >= 4
with pytest.raises(s3_dataset.S3DownloadCancelled):
s3_dataset.prepare_s3_dataset_download(
_cfg(),
cancel_callback = cancel_after_download_starts,
)
assert not target_dir.exists()
# ── S3Config model (camelCase aliases + credential validation) ──
def test_s3config_accepts_camelcase_aliases():
cfg = S3Config.model_validate(
{
"bucket": "b",
"region": "eu-west-1",
"accessKeyId": "AKIA",
"secretAccessKey": "shh",
}
)
assert cfg.access_key_id == "AKIA"
assert cfg.secret_access_key == "shh"
# model_dump() yields snake_case for the loader.
assert cfg.model_dump()["access_key_id"] == "AKIA"
def test_s3config_accepts_snake_case():
cfg = S3Config.model_validate(
{"bucket": "b", "access_key_id": "AKIA", "secret_access_key": "shh"}
)
assert cfg.access_key_id == "AKIA"
def test_s3config_requires_credentials_or_iam():
with pytest.raises(ValueError):
S3Config.model_validate({"bucket": "b"})
def test_s3config_iam_role_needs_no_keys():
cfg = S3Config.model_validate({"bucket": "b", "useIamRole": True})
assert cfg.use_iam_role is True

49
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -28,6 +28,8 @@ from core.inference.tool_call_parser import (
parse_tool_calls_from_text,
strip_tool_markup,
)
from state import tool_approvals
from state.tool_approvals import resolve_tool_decision
from utils.datasets import is_gpt_oss_model_name
@ -84,8 +86,7 @@ class TestParser:
# A code parameter with a literal </parameter> must not truncate: the
# parser uses end-of-body as the only boundary for single-param calls.
text = (
"<function=python><parameter=code>html = '<a></a>'\n"
"print('hi')</parameter></function>"
"<function=python><parameter=code>html = '<a></a>'\nprint('hi')</parameter></function>"
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
@ -1033,6 +1034,50 @@ class TestGuardrails:
_collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "x"})]
def test_confirm_tool_calls_close_after_prompt_cleans_slot(self, monkeypatch):
approval_id = "approval-close-sf"
monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: approval_id)
loop, exec_fn = _make_loop(
turns = [['<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>']],
exec_results = ["OK"],
confirm_tool_calls = True,
session_id = "sess",
max_tool_iterations = 1,
)
with tool_approvals._lock:
tool_approvals._pending.clear()
try:
assert next(loop)["type"] == "status"
start = next(loop)
assert start["type"] == "tool_start"
assert start["approval_id"] == approval_id
with tool_approvals._lock:
assert approval_id in tool_approvals._pending
finally:
loop.close()
with tool_approvals._lock:
assert approval_id not in tool_approvals._pending
assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False
assert exec_fn.calls == []
def test_confirm_tool_calls_skips_rag_autoinject(self, monkeypatch):
def fail_autoinject(*_args, **_kwargs):
raise AssertionError("RAG autoinject must not run before approval")
monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject)
loop, exec_fn = _make_loop(
turns = [["plain answer"]],
confirm_tool_calls = True,
rag_scope = {"thread_id": "t1"},
)
events = _collect_events(loop)
assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events)
assert exec_fn.calls == []
def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self):
turns = iter(
[

24
studio/backend/tests/test_startup_banner_loopback.py Normal file
View file

@ -0,0 +1,24 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression for PR #6295: the banner's canned http://127.0.0.1 URL is valid
only for the exact loopback aliases, so any other bind (e.g. a specific LAN IP)
must show its real address."""
import pytest
from startup_banner import print_studio_access_banner
def test_non_alias_loopback_shows_real_address(capsys):
# A server bound to 127.0.0.2 does not listen on 127.0.0.1.
print_studio_access_banner(port = 8891, bind_host = "127.0.0.2", display_host = "127.0.0.2")
out = capsys.readouterr().out
assert "http://127.0.0.2:8891" in out
assert "http://127.0.0.1" not in out
@pytest.mark.parametrize("host", ["127.0.0.1", "localhost"])
def test_alias_loopback_shows_canned_url(capsys, host):
print_studio_access_banner(port = 8891, bind_host = host, display_host = host)
assert "http://127.0.0.1:8891" in capsys.readouterr().out

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

@ -149,6 +149,7 @@ def test_help_output():
"--host",
"--frontend",
"--silent",
"--tensor-parallel",
]:
assert flag in out, f"Missing flag {flag!r} in --help output"
print(" PASS --help shows all flags")

578
studio/backend/tests/test_tensor_parallel.py Normal file
View file

@ -0,0 +1,578 @@
# 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 contract for the Tensor Parallelism toggle.
The toggle threads a single ``tensor_parallel`` bool from the chat UI
through the load request to a ``--split-mode tensor`` llama-server flag,
and round-trips it back via the load/status responses so the switch
reflects what is actually running. These tests pin:
* the pydantic request/response/status contract (snake_case key,
default False),
* the backend ``tensor_parallel`` property and its reset on unload,
* the ``_already_in_target_state`` reload-detection branch, and
* that ``--split-mode tensor`` is emitted only behind the toggle.
"""
from __future__ import annotations
import asyncio
import inspect
import sys
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Same external-dep stubs as the other llama_cpp unit tests so importing
# the backend doesn't drag in structlog / httpx / loggers.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
sys.modules.setdefault("structlog", _structlog_stub)
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
_httpx_stub.Client = type(
"C",
(),
{
"__init__": lambda s, **kw: None,
"__enter__": lambda s: s,
"__exit__": lambda s, *a: None,
},
)
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference import llama_cpp as llama_cpp_module
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_server_args import resolve_tensor_parallel
from core.inference.tensor_fallback import load_with_tensor_fallback
from models.inference import (
InferenceStatusResponse,
LoadRequest,
LoadResponse,
)
# ── Pydantic contract (snake_case key, default False) ────────────────
def test_load_request_defaults_tensor_parallel_false():
req = LoadRequest(model_path = "owner/repo")
assert req.tensor_parallel is False
def test_load_request_accepts_tensor_parallel():
req = LoadRequest(model_path = "owner/repo", tensor_parallel = True)
assert req.tensor_parallel is True
def test_load_request_round_trips_json_key():
# The frontend sends the snake_case key verbatim.
req = LoadRequest.model_validate({"model_path": "owner/repo", "tensor_parallel": True})
assert req.tensor_parallel is True
assert req.model_dump()["tensor_parallel"] is True
@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
def test_response_models_emit_tensor_parallel(model_cls):
# Default False, and the key is always present in the JSON body.
if model_cls is LoadResponse:
default = model_cls(
status = "loaded",
model = "owner/repo",
display_name = "repo",
inference = {},
)
on = model_cls(
status = "loaded",
model = "owner/repo",
display_name = "repo",
inference = {},
tensor_parallel = True,
)
else:
default = model_cls()
on = model_cls(tensor_parallel = True)
assert default.model_dump()["tensor_parallel"] is False
assert on.model_dump()["tensor_parallel"] is True
# ── Backend property + reset ─────────────────────────────────────────
class _FakeProcess:
"""Stand-in for subprocess.Popen so _kill_process is a no-op."""
def terminate(self):
pass
def wait(self, timeout = None):
return 0
def kill(self):
pass
def poll(self):
return 0
def test_tensor_parallel_property_defaults_false():
assert LlamaCppBackend().tensor_parallel is False
def test_tensor_parallel_property_reflects_field():
backend = LlamaCppBackend()
backend._tensor_parallel = True
assert backend.tensor_parallel is True
def test_unload_resets_tensor_parallel():
backend = LlamaCppBackend()
backend._process = _FakeProcess()
backend._tensor_parallel = True
backend.unload_model()
assert backend.tensor_parallel is False
# ── _already_in_target_state reload-detection branch ─────────────────
def _loaded_backend(tensor_parallel: bool) -> LlamaCppBackend:
backend = LlamaCppBackend()
backend._process = _FakeProcess() # is_loaded only checks "is not None"
backend._healthy = True
backend._model_identifier = "owner/repo"
backend._hf_variant = "Q4_K_M"
backend._requested_n_ctx = 8192
backend._cache_type_kv = None
backend._requested_spec_mode = "auto"
backend._chat_template_override = None
backend._is_vision = False
backend._extra_args = None
backend._gguf_path = None
backend._tensor_parallel = tensor_parallel
return backend
def _target_state(backend: LlamaCppBackend, tensor_parallel: bool) -> bool:
return backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "auto",
chat_template_override = None,
extra_args = None,
is_vision = False,
tensor_parallel = tensor_parallel,
)
@pytest.mark.parametrize("flag", [True, False])
def test_already_in_target_state_matches_same_tensor_parallel(flag):
assert _target_state(_loaded_backend(flag), flag) is True
@pytest.mark.parametrize(
"loaded,requested",
[(False, True), (True, False)],
)
def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, requested):
# Flipping the toggle either direction must force a reload so the
# command is rebuilt with/without --split-mode tensor.
assert _target_state(_loaded_backend(loaded), requested) is False
def test_already_in_target_state_reconciles_split_mode_extras():
# Tensor engaged via --split-mode in extras (boolean omitted/default False)
# must match a server already running tensor mode -- no spurious reload.
backend = _loaded_backend(tensor_parallel = True)
backend._extra_args = ["--split-mode", "tensor"]
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "auto",
chat_template_override = None,
extra_args = ["--split-mode", "tensor"],
is_vision = False,
tensor_parallel = False,
)
is True
)
# ── --split-mode tensor is emitted only behind the toggle ────────────
def _load_model_source() -> str:
return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)
def test_split_mode_tensor_is_gated_on_the_toggle():
src = _load_model_source()
assert (
'cmd.extend(["--split-mode", "tensor"])' in src
), "the tensor-parallel flag emission must be present in load_model"
# The emission lives behind `if tensor_parallel:` -- it must never be
# part of the unconditional base cmd list.
base_start = src.find("cmd = [")
base_end = src.find("\n ]", base_start)
base_block = src[base_start:base_end] if base_end > base_start else ""
assert (
"--split-mode" not in base_block
), "--split-mode must be conditional, not in the base cmd list"
gate = src.find("if tensor_parallel:")
emit = src.find('cmd.extend(["--split-mode", "tensor"])')
assert 0 <= gate < emit, "emission must sit under `if tensor_parallel:`"
def test_proportional_tensor_split_is_emitted_in_tensor_mode():
# Asymmetric GPUs (e.g. 48 GB + 24 GB) OOM the smaller card under the
# even default; the allocator weights --tensor-split by free VRAM. Pin
# that the flag is emitted from inside the tensor-parallel block.
src = _load_model_source()
assert '"--tensor-split"' in src
gate = src.find("if tensor_parallel:")
ts = src.find('"--tensor-split"')
nxt_else = src.find("self._tensor_parallel = False")
assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`"
# ── tensor-mode allocation: conservative VRAM budget ─────────────────
def _kv_seeded_backend() -> LlamaCppBackend:
# Minimal GGUF metadata so _can_estimate_kv() is True (legacy KV path).
backend = LlamaCppBackend()
backend._n_layers = 32
backend._embedding_length = 4096
backend._n_heads = 32
backend._n_kv_heads = 8
backend._context_length = 131072
return backend
def test_fit_context_budget_frac_override_is_tighter():
backend = _kv_seeded_backend()
model_size = 8 * 1024**3
pool_mib = 24 * 1024 # tight enough that KV capping bites
fit_default = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16")
fit_tp = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16", budget_frac = 0.80)
assert fit_tp < 131072, "expected the context to be capped at this VRAM tier"
assert fit_tp <= fit_default, "a tighter budget must not allow MORE context"
# Omitting the override must reproduce the default budget exactly.
assert backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") == fit_default
# ── unsupported-arch load failure -> clean message ───────────────────
def test_split_mode_tensor_arch_failure_message():
msg = LlamaCppBackend._classify_llama_start_failure(
"llama_model_create: LLAMA_SPLIT_MODE_TENSOR not implemented for "
"architecture 'deepseek2'",
None,
"unsloth/DeepSeek-V3-GGUF",
)
assert "Tensor parallelism is not supported" in msg
def test_unrelated_arch_failure_not_hijacked_by_tensor_message():
msg = LlamaCppBackend._classify_llama_start_failure(
"unknown model architecture: 'flux'", "/models/flux.gguf", None
)
assert "Tensor parallelism" not in msg
# ── _plan_tensor_parallel: the allocation math (pure, no model/GPU) ───
# Seeded full-attention KV (~128 KiB/token) via _kv_seeded_backend, so the
# context cap + split are deterministic. Asserts relationships rather than
# magic numbers so the KV estimate can evolve without breaking these.
_GB = 1024**3
_ASYM = [(0, 48000), (1, 24000)] # asymmetric pool, 72000 MiB
_SYM = [(0, 24000), (1, 24000)] # symmetric pool
def _plan(
model_gb,
target = 131072,
gpus = _ASYM,
mtp = False,
):
b = _kv_seeded_backend()
return b, b._plan_tensor_parallel(gpus, int(model_gb * _GB), target, mtp_engaged = mtp)
def _kv_budget_b(model_gb, gpus = _ASYM):
reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB)
def test_tp_plan_weighted_split_on_asymmetric_big_model():
b, (ec, mac, gi, ts) = _plan(50)
reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
assert gi == [0, 1]
# split weighted by (free - buffer), not raw free
assert ts == [48000 - reserve, 24000 - reserve]
assert ec < 131072 # capped below native
def test_tp_plan_even_split_when_model_fits():
# A small model whose even share fits the smallest GPU -> llama.cpp's even
# default (None), which is safe for archs that crash on a weighted split.
_, (ec, mac, gi, ts) = _plan(4)
assert ts is None
def test_tp_plan_symmetric_gpus_use_even_split():
_, (ec, mac, gi, ts) = _plan(8, gpus = _SYM)
assert ts is None
def test_tp_plan_context_fits_pool_budget_no_oom():
b, (ec, mac, gi, ts) = _plan(50)
# the chosen context's KV must fit the pooled budget (weights + buffers)
assert b._estimate_kv_cache_bytes(ec) <= _kv_budget_b(50)
def test_tp_plan_uses_available_vram_not_wasteful():
# when the cap engages, the chosen context nearly fills the budget
b, (ec, mac, gi, ts) = _plan(50)
assert b._estimate_kv_cache_bytes(ec) >= 0.9 * _kv_budget_b(50)
def test_tp_plan_weights_exceed_pool_floors_context():
# 70 GB > pool minus per-GPU reserves -> floor (triggers layer fallback)
_, (ec, mac, gi, ts) = _plan(70)
assert ec == 2048
def test_tp_plan_floor_never_exceeds_explicit_small_context():
# An explicit context below the 2048 floor must not be raised: a caller
# asking for 1024 should not have KV sized for 2048 (avoidable OOM).
_, (ec, mac, gi, ts) = _plan(70, target = 1024) # weights exceed pool -> floor path
assert ec == 1024
_, (ec2, *_rest) = _plan(50, target = 1024) # cap path with a tiny budget
assert ec2 <= 1024
def test_tp_plan_explicit_context_honored_when_it_fits():
_, (ec, mac, gi, ts) = _plan(50, target = 8192)
assert ec == 8192
def test_tp_plan_explicit_context_capped_when_too_large():
_, (ec, mac, gi, ts) = _plan(50, target = 131072)
assert 2048 <= ec < 131072
def test_tp_plan_max_available_ctx_reports_native_not_explicit_ctx():
# An explicit small ctx caps effective_ctx but the UI ceiling
# (max_available_ctx) must reflect the native/hardware cap, not the request.
b = _kv_seeded_backend()
ec, mac, _gi, _ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 8192, max_target_ctx = 131072)
_, native_mac, *_ = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072)
assert ec == 8192 # explicit request honored for the load
assert mac == native_mac > ec # ceiling reflects the hardware cap
def test_tp_plan_mtp_reserves_extra_and_shrinks_context():
_, (ec_no, *_rest) = _plan(50)
_, (ec_mtp, *_rest) = _plan(50, mtp = True)
assert ec_mtp < ec_no
def test_tp_plan_no_kv_metadata_floors_context():
b = LlamaCppBackend() # no KV metadata -> can't size safely
ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072)
assert ec <= 4096
def test_tp_plan_single_gpu_never_splits():
# The toggle is a no-op without >= 2 GPUs (most dev/CI machines). Even if
# the planner is reached, it must not emit a tensor split.
b = _kv_seeded_backend()
ec, mac, gi, ts = b._plan_tensor_parallel([(0, 24000)], int(8 * _GB), 8192)
assert ts is None
assert gi == [0]
def test_tp_plan_zero_gpus_never_splits():
b = _kv_seeded_backend()
ec, mac, gi, ts = b._plan_tensor_parallel([], int(8 * _GB), 8192)
assert ts is None
assert gi == []
def test_tp_plan_drops_gpu_below_buffer_reserve():
# A GPU with less free VRAM than the per-device compute-buffer reserve
# can't host tensor mode; it's excluded, which here leaves <2 usable -> no
# split (and gpu_indices reflects only the usable device).
b = _kv_seeded_backend()
reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, reserve - 1)], int(8 * _GB), 8192)
assert gi == [0]
assert ts is None
# ── route auto-fallback survives a *raised* tensor-load crash ─────────
# A tensor-incompatible model makes load_model RAISE (Gemma 3n aborts) rather
# than return False. The /load fallback helper must catch that and retry with
# layer split -- stripping any --split-mode from the extras so the retry can't
# relaunch tensor -- while a non-tensor load propagates its exception. These
# exercise the real helper with a fake loader (no GPU, no llama-server).
class _RecordingLoader:
"""Fake ``attempt_load``: crashes whenever tensor mode is effectively
engaged (via the bool or a ``--split-mode`` in extras), like a real
tensor-incompatible model; succeeds on layer split."""
def __init__(self):
self.calls: list[tuple] = []
async def __call__(self, tensor_parallel, extra_args):
self.calls.append((tensor_parallel, list(extra_args) if extra_args else extra_args))
if resolve_tensor_parallel(extra_args, tensor_parallel):
raise RuntimeError("llama-server failed to start")
return True
def test_tensor_fallback_retries_layer_on_crash():
loader = _RecordingLoader()
ok = asyncio.run(
load_with_tensor_fallback(loader, requested_tensor = True, extra_args = None, label = "m")
)
assert ok is True
# tensor first (crashes), then layer split.
assert [c[0] for c in loader.calls] == [True, False]
def test_tensor_fallback_no_retry_on_success():
calls: list[bool] = []
async def _ok(tensor_parallel, extra_args):
calls.append(tensor_parallel)
return True
ok = asyncio.run(
load_with_tensor_fallback(_ok, requested_tensor = True, extra_args = None, label = "m")
)
assert ok is True
assert calls == [True] # no fallback when the tensor load succeeds
def test_tensor_fallback_retries_when_tensor_returns_false():
# load_model can signal failure by *returning False* (not only by raising);
# that must trigger the layer-split retry just like a crash does.
calls: list[bool] = []
async def _false_on_tensor(tensor_parallel, extra_args):
calls.append(tensor_parallel)
return not resolve_tensor_parallel(extra_args, tensor_parallel)
ok = asyncio.run(
load_with_tensor_fallback(
_false_on_tensor, requested_tensor = True, extra_args = None, label = "m"
)
)
assert ok is True
assert calls == [True, False]
def test_tensor_fallback_returns_false_when_both_attempts_fail():
# Tensor fails and the layer retry also fails -> the helper returns False so
# the route raises its own HTTP 500 (it does not crash mid-flight).
calls: list[bool] = []
async def _always_false(tensor_parallel, extra_args):
calls.append(tensor_parallel)
return False
ok = asyncio.run(
load_with_tensor_fallback(_always_false, requested_tensor = True, extra_args = None, label = "m")
)
assert ok is False
assert calls == [True, False] # tried tensor, then layer split
def test_tensor_fallback_skips_layer_retry_when_cancelled():
# load_model returns False on a user cancellation too. When cancelled() is
# True, the helper must NOT relaunch the load the user just cancelled.
calls: list[bool] = []
async def _false_on_tensor(tensor_parallel, extra_args):
calls.append(tensor_parallel)
return False
ok = asyncio.run(
load_with_tensor_fallback(
_false_on_tensor,
requested_tensor = True,
extra_args = None,
label = "m",
cancelled = lambda: True,
)
)
assert ok is False
assert calls == [True] # no layer-split retry after cancellation
@pytest.mark.parametrize(
"extras",
[
["--split-mode", "tensor", "-c", "4096"],
["-sm", "tensor", "-c", "4096"],
["--split-mode=tensor", "-c", "4096"],
["-sm=tensor", "-c", "4096"],
],
)
def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras):
# Tensor engaged via extras (boolean False); the retry must drop every
# --split-mode form (long/short, space/=) but keep the user's other flags,
# else resolve_tensor_parallel re-enables tensor and relaunches the crash.
loader = _RecordingLoader()
ok = asyncio.run(
load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m")
)
assert ok is True
assert len(loader.calls) == 2
assert loader.calls[1][1] == ["-c", "4096"] # split-mode stripped, -c kept
def test_tensor_fallback_propagates_non_tensor_crash():
async def _always_raise(tensor_parallel, extra_args):
raise RuntimeError("bad model")
with pytest.raises(RuntimeError, match = "bad model"):
asyncio.run(
load_with_tensor_fallback(
_always_raise, requested_tensor = False, extra_args = None, label = "m"
)
)

261
studio/backend/tests/test_tool_approvals.py Normal file
View file

@ -0,0 +1,261 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Concurrency tests for the per-call tool-call confirmation gate.
``state.tool_approvals`` coordinates two threads: the agentic loop thread
blocked in ``wait_tool_decision`` and the request thread that delivers the
user's choice through ``resolve_tool_decision``. Each gated call carries a
unique ``approval_id`` so a stale or concurrent confirmation can never
resolve the wrong call. These tests exercise that handshake directly --
no model, no server -- so the race windows are fast and deterministic.
"""
import threading
import time
import pytest
from state import tool_approvals
from state.tool_approvals import (
TOOL_REJECTED_MESSAGE,
abort_tool_decision,
begin_tool_decision,
new_approval_id,
request_tool_decision,
resolve_tool_decision,
wait_tool_decision,
)
@pytest.fixture(autouse = True)
def _clear_pending():
"""Each test starts and ends with an empty ``_pending`` map."""
with tool_approvals._lock:
tool_approvals._pending.clear()
yield
with tool_approvals._lock:
tool_approvals._pending.clear()
class _Waiter:
"""Run ``request_tool_decision`` in a thread and capture its result."""
def __init__(
self,
session_id,
approval_id,
cancel_event = None,
timeout = None,
):
self.session_id = session_id
self.approval_id = approval_id
self.cancel_event = cancel_event
self.timeout = timeout
self.result = None
self._thread = threading.Thread(target = self._run, daemon = True)
def _run(self):
kwargs = {"cancel_event": self.cancel_event}
if self.timeout is not None:
kwargs["timeout"] = self.timeout
self.result = request_tool_decision(self.session_id, self.approval_id, **kwargs)
def start(self):
self._thread.start()
_wait_until(lambda: _has_pending(self.approval_id))
return self
def join(self, timeout = 5.0):
self._thread.join(timeout = timeout)
assert not self._thread.is_alive(), "waiter thread did not finish"
return self.result
def _has_pending(approval_id) -> bool:
with tool_approvals._lock:
return approval_id in tool_approvals._pending
def _wait_until(
pred,
timeout = 2.0,
interval = 0.005,
) -> bool:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if pred():
return True
time.sleep(interval)
return False
# ── Basic allow / deny ───────────────────────────────────────────────
def test_allow_decision():
aid = new_approval_id()
w = _Waiter("sess", aid).start()
assert resolve_tool_decision(aid, "allow", session_id = "sess") is True
assert w.join() == "allow"
def test_deny_decision():
aid = new_approval_id()
w = _Waiter("sess", aid).start()
assert resolve_tool_decision(aid, "deny", session_id = "sess") is True
assert w.join() == "deny"
def test_slot_cleaned_up_after_decision():
aid = new_approval_id()
w = _Waiter("sess", aid).start()
resolve_tool_decision(aid, "allow")
w.join()
assert _wait_until(lambda: not _has_pending(aid))
def test_abort_tool_decision_removes_unwaited_slot():
aid = new_approval_id()
slot = begin_tool_decision("sess", aid)
abort_tool_decision(slot, aid)
assert not _has_pending(aid)
assert resolve_tool_decision(aid, "allow", session_id = "sess") is False
def test_approval_ids_are_unique():
ids = {new_approval_id() for _ in range(1000)}
assert len(ids) == 1000
# ── Pre-registration race (begin before wait) ────────────────────────
def test_resolve_before_wait_is_not_lost():
"""A decision delivered after ``begin`` but before ``wait`` survives.
The loop registers the slot before it yields ``tool_start``, so even a
confirmation that races ahead of the blocking ``wait`` is recorded on
the slot and returned -- never dropped.
"""
aid = new_approval_id()
slot = begin_tool_decision("sess", aid)
assert resolve_tool_decision(aid, "allow", session_id = "sess") is True
# wait() is only entered now, after the decision already landed.
assert wait_tool_decision(slot, aid) == "allow"
assert not _has_pending(aid)
# ── Resolver edge cases ──────────────────────────────────────────────
def test_resolve_unknown_approval_returns_false():
assert resolve_tool_decision(new_approval_id(), "allow") is False
def test_resolve_empty_approval_returns_false():
assert resolve_tool_decision("", "allow") is False
assert resolve_tool_decision(None, "allow") is False
def test_resolve_wrong_session_scope_returns_false():
aid = new_approval_id()
w = _Waiter("sess-a", aid).start()
# Correct approval_id but the wrong session must not resolve it.
assert resolve_tool_decision(aid, "allow", session_id = "sess-b") is False
assert _has_pending(aid)
# The right session still works.
assert resolve_tool_decision(aid, "allow", session_id = "sess-a") is True
assert w.join() == "allow"
def test_duplicate_resolve_after_completion_returns_false():
aid = new_approval_id()
w = _Waiter("sess", aid).start()
assert resolve_tool_decision(aid, "allow") is True
w.join()
assert _wait_until(lambda: not _has_pending(aid))
assert resolve_tool_decision(aid, "deny") is False
def test_first_decision_is_immutable():
"""A second confirmation cannot flip an already-recorded decision.
The waiter reads ``slot["decision"]`` outside the lock and then cleans up,
so a duplicate or out-of-order POST that lands in that window must be
rejected and must not overwrite the first decision -- an Allow can never
become a Deny. Distinct from the after-completion case above: here the slot
is still pending (no waiter has consumed it yet).
"""
aid = new_approval_id()
slot = begin_tool_decision("sess", aid)
assert resolve_tool_decision(aid, "allow", session_id = "sess") is True
# Second decision, same id, before any waiter consumes/cleans the slot.
assert resolve_tool_decision(aid, "deny", session_id = "sess") is False
assert slot["decision"] == "allow"
# The waiter still observes the first (immutable) decision.
assert wait_tool_decision(slot, aid) == "allow"
assert not _has_pending(aid)
# ── Cancellation and timeout ─────────────────────────────────────────
def test_cancel_event_breaks_wait_as_deny():
cancel = threading.Event()
aid = new_approval_id()
w = _Waiter("sess", aid, cancel_event = cancel).start()
cancel.set()
assert w.join(timeout = 3.0) == "deny"
assert _wait_until(lambda: not _has_pending(aid))
def test_timeout_returns_deny():
aid = new_approval_id()
start = time.monotonic()
result = request_tool_decision("sess", aid, timeout = 0.1)
assert result == "deny"
assert time.monotonic() - start < 2.0
assert not _has_pending(aid)
# ── Independence across concurrent calls ─────────────────────────────
def test_two_pending_calls_same_session_are_independent():
"""Keying on approval_id, not session, keeps concurrent calls distinct.
Resolving the first call's id must not unblock or alter the second
call pending in the same session.
"""
a1, a2 = new_approval_id(), new_approval_id()
w1 = _Waiter("sess", a1).start()
w2 = _Waiter("sess", a2).start()
assert resolve_tool_decision(a1, "deny", session_id = "sess") is True
assert w1.join() == "deny"
# w2 is still waiting on its own id.
assert _has_pending(a2)
assert resolve_tool_decision(a2, "allow", session_id = "sess") is True
assert w2.join() == "allow"
def test_concurrent_distinct_calls_route_their_own_decisions():
n = 25
waiters = {}
for i in range(n):
aid = new_approval_id()
waiters[aid] = _Waiter(f"s{i}", aid).start()
expected = {aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters)}
for aid, decision in expected.items():
assert resolve_tool_decision(aid, decision) is True
for aid, w in waiters.items():
assert w.join() == expected[aid]
# ── Constants ────────────────────────────────────────────────────────
def test_rejected_message_is_user_facing_text():
assert isinstance(TOOL_REJECTED_MESSAGE, str)
assert TOOL_REJECTED_MESSAGE.strip()

170
studio/backend/tests/test_tool_confirm_loop.py Normal file
View file

@ -0,0 +1,170 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Integration tests for the confirmation gate inside the real tool loop.
These drive ``run_safetensors_tool_loop`` (no model -- hand-crafted fake
generators) with ``confirm_tool_calls=True`` and resolve each pending
decision inline. The slot is registered before ``tool_start`` is yielded,
so resolving right after receiving that event always lands before the
loop blocks. Covers: allow executes once, deny skips execution and feeds
back the rejection, disabled/duplicate calls are not prompted, and a
denied call does not pollute duplicate detection.
"""
import pytest
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from state import tool_approvals
from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision
_SESSION = "loop-session"
@pytest.fixture(autouse = True)
def _clear_pending():
with tool_approvals._lock:
tool_approvals._pending.clear()
yield
with tool_approvals._lock:
tool_approvals._pending.clear()
class _FakeExecuteTool:
def __init__(self):
self.calls = []
def __call__(
self,
name,
arguments,
*,
cancel_event = None,
timeout = None,
session_id = None,
rag_scope = None,
):
self.calls.append((name, arguments))
return f"RESULT[{name}]"
def _tool_call(name, args_json):
return f'<tool_call>{{"name": "{name}", "arguments": {args_json}}}</tool_call>'
def _multi_turn(turns):
"""A single_turn generator that yields one full snapshot per turn."""
turn_iter = iter(turns)
def _gen(_messages):
try:
yield next(turn_iter)
except StopIteration:
return
return _gen
_DEFAULT_TOOLS = [
{"type": "function", "function": {"name": "python"}},
{"type": "function", "function": {"name": "web_search"}},
]
def _drive(
turns,
decisions,
*,
tools = None,
):
"""Run the loop, resolving each gated tool_start with the next decision.
The advertised ``tools`` list drives the loop's enabled-tool filter
(pass a list omitting a tool to make a call to it "disabled").
Returns (events, execute_calls).
"""
decision_iter = iter(decisions)
exec_fn = _FakeExecuteTool()
gen = run_safetensors_tool_loop(
single_turn = _multi_turn(turns),
messages = [{"role": "user", "content": "hi"}],
tools = _DEFAULT_TOOLS if tools is None else tools,
execute_tool = exec_fn,
session_id = _SESSION,
confirm_tool_calls = True,
)
events = []
for ev in gen:
events.append(ev)
if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"):
# Slot is already registered (begin ran before this yield), so
# the decision lands before the loop enters its blocking wait.
resolve_tool_decision(ev["approval_id"], next(decision_iter), session_id = _SESSION)
return events, exec_fn.calls
def _tool_starts(events):
return [e for e in events if e["type"] == "tool_start"]
def _tool_ends(events):
return [e for e in events if e["type"] == "tool_end"]
def test_allow_executes_the_tool_once():
events, calls = _drive(
[_tool_call("python", '{"code": "print(1)"}'), "final answer"],
["allow"],
)
starts = _tool_starts(events)
assert len(starts) == 1
assert starts[0]["awaiting_confirmation"] is True
assert starts[0]["approval_id"]
assert calls == [("python", {"code": "print(1)"})]
assert _tool_ends(events)[0]["result"] == "RESULT[python]"
def test_deny_skips_execution_and_feeds_rejection():
events, calls = _drive(
[_tool_call("python", '{"code": "print(1)"}'), "final answer"],
["deny"],
)
assert calls == [] # tool never ran
assert _tool_ends(events)[0]["result"] == TOOL_REJECTED_MESSAGE
def test_disabled_tool_is_not_prompted():
events, calls = _drive(
[_tool_call("python", '{"code": "print(1)"}'), "final answer"],
[],
tools = [{"type": "function", "function": {"name": "web_search"}}],
)
assert _tool_starts(events) == []
assert _tool_ends(events) == []
assert calls == []
def test_duplicate_call_is_not_prompted():
same = _tool_call("python", '{"code": "print(1)"}')
events, calls = _drive([same, same, "final answer"], ["allow"])
starts = _tool_starts(events)
assert len(starts) == 1
assert starts[0]["awaiting_confirmation"] is True
assert calls == [("python", {"code": "print(1)"})]
assert len(_tool_ends(events)) == 1
def test_denied_call_can_be_reissued_and_approved():
# Deny, then the model re-issues the identical call -> approving it must
# execute, not get suppressed as a duplicate (denied calls are not added
# to the duplicate-detection history).
same = _tool_call("python", '{"code": "print(1)"}')
events, calls = _drive([same, same, "final answer"], ["deny", "allow"])
starts = _tool_starts(events)
assert len(starts) == 2
assert starts[0]["awaiting_confirmation"] is True
assert starts[1]["awaiting_confirmation"] is True # not treated as dup
assert calls == [("python", {"code": "print(1)"})] # ran once, on approve
ends = _tool_ends(events)
assert ends[0]["result"] == TOOL_REJECTED_MESSAGE
assert ends[1]["result"] == "RESULT[python]"

219
studio/backend/tests/test_tool_confirm_stream.py Normal file
View file

@ -0,0 +1,219 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""End-to-end handshake test for the tool-confirmation gate, no model.
The real Studio stream wrappers in ``routes/inference.py`` drive the
synchronous agentic generator with ``await asyncio.to_thread(next, gen,
...)`` so the blocking ``threading.Event`` wait runs off the event loop.
This test rebuilds that exact pattern around the real
``state.tool_approvals`` functions, served by a real uvicorn process on
loopback (the same server Studio uses), and proves the load-bearing
property:
* ``tool_start`` reaches the client before the gate blocks, and
* the separate ``/tool-confirm`` POST is served *while* the stream
connection is blocked, after which the stream resumes with the executed
(allow) or rejected (deny) result -- i.e. no deadlock.
Each scenario runs under a socket-level timeout, so a regression that
reintroduces a deadlock fails fast instead of hanging the suite.
"""
import asyncio
import json
import socket
import threading
import time
import httpx
import pytest
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from state import tool_approvals
from state.tool_approvals import (
TOOL_REJECTED_MESSAGE,
begin_tool_decision,
new_approval_id,
resolve_tool_decision,
wait_tool_decision,
)
_EXECUTED_RESULT = "tool executed: 2"
@pytest.fixture(autouse = True)
def _clear_pending():
with tool_approvals._lock:
tool_approvals._pending.clear()
yield
with tool_approvals._lock:
tool_approvals._pending.clear()
def _build_app() -> FastAPI:
"""Minimal app mirroring the real stream/confirm wiring."""
app = FastAPI()
def agentic_gen(session_id, cancel_event):
# Same shape as the real loops: register the approval slot, announce
# the call (echoing approval_id), gate on the decision, then either
# execute or feed back the rejection.
approval_id = new_approval_id()
slot = begin_tool_decision(session_id, approval_id)
yield {
"type": "tool_start",
"tool_name": "python",
"approval_id": approval_id,
"awaiting_confirmation": True,
}
denied = wait_tool_decision(slot, approval_id, cancel_event = cancel_event) == "deny"
result = TOOL_REJECTED_MESSAGE if denied else _EXECUTED_RESULT
yield {"type": "tool_end", "tool_name": "python", "result": result}
@app.post("/stream")
async def stream(req: Request):
body = await req.json()
session_id = body.get("session_id")
cancel_event = threading.Event()
sentinel = object()
async def wrapper():
gen = agentic_gen(session_id, cancel_event)
while True:
event = await asyncio.to_thread(next, gen, sentinel)
if event is sentinel:
break
yield f"data: {json.dumps(event)}\n\n"
return StreamingResponse(wrapper(), media_type = "text/event-stream")
@app.post("/tool-confirm")
async def tool_confirm(req: Request):
body = await req.json()
resolved = resolve_tool_decision(
body.get("approval_id"),
body.get("decision"),
session_id = body.get("session_id"),
)
return {"resolved": resolved}
return app
def _free_port() -> int:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
s.close()
return port
class _Server:
"""Run a uvicorn server in a background thread for the test's lifetime."""
def __init__(self, app):
self.port = _free_port()
config = uvicorn.Config(app, host = "127.0.0.1", port = self.port, log_level = "warning")
self.server = uvicorn.Server(config)
self._thread = threading.Thread(target = self.server.run, daemon = True)
def __enter__(self):
self._thread.start()
deadline = time.monotonic() + 10.0
while time.monotonic() < deadline:
if self.server.started:
return self
time.sleep(0.02)
raise AssertionError("uvicorn did not start in time")
def __exit__(self, *exc):
self.server.should_exit = True
self._thread.join(timeout = 10.0)
@property
def base_url(self) -> str:
return f"http://127.0.0.1:{self.port}"
async def _gate_is_blocking(approval_id) -> None:
"""Wait until the stream thread is parked on this approval's slot.
The slot is registered before ``tool_start`` is yielded, so it exists
by the time the client receives the event -- exactly as in reality,
where the confirm POST only arrives after the card renders.
"""
for _ in range(400):
with tool_approvals._lock:
slot = tool_approvals._pending.get(approval_id)
if slot is not None and not slot["event"].is_set():
return
await asyncio.sleep(0.005)
raise AssertionError("gate never started waiting")
async def _drive(base_url, session_id, decision):
events = []
resolved = None
timeout = httpx.Timeout(10.0)
async with httpx.AsyncClient(base_url = base_url, timeout = timeout) as client:
async with client.stream("POST", "/stream", json = {"session_id": session_id}) as resp:
assert resp.status_code == 200
async for line in resp.aiter_lines():
if not line.startswith("data: "):
continue
event = json.loads(line[len("data: ") :])
events.append(event)
if event["type"] == "tool_start":
# The stream is now blocked on the gate; the confirm
# POST (echoing approval_id) must still be served over a
# second connection.
approval_id = event["approval_id"]
await _gate_is_blocking(approval_id)
r = await client.post(
"/tool-confirm",
json = {
"session_id": session_id,
"approval_id": approval_id,
"decision": decision,
},
)
resolved = r.json()["resolved"]
return events, resolved
def _run(session_id, decision):
with _Server(_build_app()) as srv:
return asyncio.run(
asyncio.wait_for(_drive(srv.base_url, session_id, decision), timeout = 15.0)
)
def _types(events):
return [e["type"] for e in events]
def test_allow_resumes_stream_with_executed_result():
events, resolved = _run("sess-allow", "allow")
assert resolved is True
assert _types(events) == ["tool_start", "tool_end"]
assert events[-1]["result"] == _EXECUTED_RESULT
def test_deny_resumes_stream_with_rejection_result():
events, resolved = _run("sess-deny", "deny")
assert resolved is True
assert _types(events) == ["tool_start", "tool_end"]
assert events[-1]["result"] == TOOL_REJECTED_MESSAGE
def test_tool_start_precedes_the_block_and_carries_approval_id():
# The first streamed event is always tool_start, proving the buttons
# can render before the backend pauses for the decision -- and it
# carries the approval_id / awaiting_confirmation the UI needs.
events, _ = _run("sess-order", "allow")
assert events[0]["type"] == "tool_start"
assert events[0]["awaiting_confirmation"] is True
assert events[0]["approval_id"]

181
studio/backend/tests/test_training_raw_support.py
View file

@ -107,10 +107,191 @@ class TestTrainingRawSupport(unittest.TestCase):
model_name = "unsloth/test",
training_type = "LoRA/QLoRA",
max_grad_norm = 0.7,
max_grad_value = 3.0,
max_grad_leaf_norm = 1.3,
)
config = mock_process.call_args.kwargs["kwargs"]["config"]
self.assertEqual(config["max_grad_norm"], 0.7)
self.assertEqual(config["max_grad_value"], 3.0)
self.assertEqual(config["max_grad_leaf_norm"], 1.3)
def test_training_backend_forwards_random_seed_without_internal_mlx_seed_keys(self):
backend = TrainingBackend()
class DummyProcess:
pid = 12345
def start(self):
return None
class DummyThread:
def start(self):
return None
dummy_queue = object()
with (
patch(
"core.training.training.prepare_gpu_selection",
return_value = ([0], {"selection_mode": "auto"}),
),
patch(
"core.training.training._CTX.Queue",
side_effect = [dummy_queue, dummy_queue],
),
patch(
"core.training.training._CTX.Process", return_value = DummyProcess()
) as mock_process,
patch(
"core.training.training.threading.Thread",
return_value = DummyThread(),
),
):
backend.start_training(
job_id = "test-seed",
model_name = "unsloth/test",
training_type = "LoRA/QLoRA",
random_seed = 1234,
)
config = mock_process.call_args.kwargs["kwargs"]["config"]
self.assertEqual(config["random_seed"], 1234)
self.assertNotIn("model_random_state", config)
self.assertNotIn("lora_random_state", config)
def test_route_forwards_all_grad_clipping_fields(self):
# The HTTP route builds the config dict by hand; a schema field that
# is not forwarded here is silently dropped for REST callers.
source = (_BACKEND_ROOT / "routes" / "training.py").read_text()
self.assertIn('"max_grad_norm": request.max_grad_norm', source)
self.assertIn('"max_grad_value": request.max_grad_value', source)
self.assertIn('"max_grad_leaf_norm": request.max_grad_leaf_norm', source)
def test_mlx_worker_falls_back_init_seeds_to_random_seed(self):
source = (_BACKEND_ROOT / "core" / "training" / "worker.py").read_text()
# random_seed itself is normalized first so explicit None coming
# from a raw / backend caller does not propagate through the chain.
self.assertIn('_raw_seed = config.get("random_seed", 3407)', source)
self.assertIn(
"random_seed = 3407 if _raw_seed is None else int(_raw_seed)",
source,
)
# Both absent and explicit None must fall back to random_seed.
# `dict.get(key, default)` only fills the default on absent keys,
# so an explicit `None` would otherwise reach FastMLXModel /
# get_peft_model and disable deterministic init.
self.assertIn('_model_seed = config.get("model_random_state")', source)
self.assertIn(
"model_random_state = random_seed if _model_seed is None else int(_model_seed)",
source,
)
self.assertIn('_lora_seed = config.get("lora_random_state")', source)
self.assertIn(
"lora_random_state = random_seed if _lora_seed is None else int(_lora_seed)",
source,
)
self.assertIn("random_state = model_random_state", source)
self.assertIn("random_state = lora_random_state", source)
# MLXTrainingConfig now receives the normalized seed directly.
self.assertIn("seed = random_seed,", source)
def test_mlx_worker_preserves_null_max_grad_value_for_trainer_default(self):
source = (_BACKEND_ROOT / "core" / "training" / "worker.py").read_text()
# None must survive to the MLX trainer so it picks its own runtime
# default, and any other value must coerce to float without
# rebinding None to 1.0 (which the legacy code did).
self.assertIn('max_grad_value = config.get("max_grad_value")', source)
self.assertIn("max_grad_value = float(max_grad_value)", source)
self.assertNotIn(
"max_grad_value = 1.0 if max_grad_value is None else float(max_grad_value)",
source,
)
def test_training_backend_normalizes_explicit_none_seed_and_dtypes(self):
# Raw / backend callers can pass `random_seed=None`,
# `cast_norm_output_to_input_dtype=None`, and MLX clip knobs
# as None (or omit them) and must NOT leak the
# `None` past `TrainingBackend.start_training`. Otherwise
# transformers.set_seed(None) raises, PEFT init becomes
# nondeterministic, and the MLX norm-output cast silently flips.
from core.training.training import (
_coerce_seed,
_coerce_optional_bool,
_coerce_optional_nonneg_float,
)
self.assertEqual(_coerce_seed(None), 3407)
self.assertEqual(_coerce_seed("123"), 123)
self.assertEqual(_coerce_seed("not-a-number"), 3407)
self.assertTrue(_coerce_optional_bool(None, True))
self.assertFalse(_coerce_optional_bool(None, False))
self.assertFalse(_coerce_optional_bool("false", True))
self.assertTrue(_coerce_optional_bool("true", False))
self.assertIsNone(_coerce_optional_nonneg_float("max_grad_value", None))
self.assertEqual(_coerce_optional_nonneg_float("max_grad_value", "2.5"), 2.5)
self.assertEqual(_coerce_optional_nonneg_float("max_grad_value", 0), 0.0)
with self.assertRaises(ValueError):
_coerce_optional_nonneg_float("max_grad_value", -1)
self.assertIsNone(_coerce_optional_nonneg_float("max_grad_leaf_norm", None))
self.assertEqual(
_coerce_optional_nonneg_float("max_grad_leaf_norm", "1.3"),
1.3,
)
with self.assertRaises(ValueError):
_coerce_optional_nonneg_float("max_grad_leaf_norm", -1)
def test_mlx_worker_feature_detects_optional_mlx_config_fields(self):
# `cast_norm_output_to_input_dtype`, `dataset_order`,
# `max_grad_leaf_norm`, and `append_eos` ship in the paired
# unsloth-zoo update. Until that floor is in place, the
# worker must gate them so releases that predate those fields can
# still construct MLXTrainingConfig without TypeError.
source = (_BACKEND_ROOT / "core" / "training" / "worker.py").read_text()
self.assertIn(
'getattr(MLXTrainingConfig, "__dataclass_fields__", {})',
source,
)
self.assertIn('if "cast_norm_output_to_input_dtype" in _supported_fields:', source)
self.assertIn('if "dataset_order" in _supported_fields:', source)
self.assertIn('if "max_grad_leaf_norm" in _supported_fields:', source)
self.assertIn(
'mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm',
source,
)
self.assertIn('if "append_eos" in _supported_fields:', source)
self.assertIn('format_type == "raw"', source)
self.assertIn('mlx_config_kwargs["append_eos"] = bool(raw_text_mode)', source)
# The unconditional kwargs must NOT include any gated field.
# Use proper paren tracking; `source.find(")", ...)` would stop at
# the first close paren inside the dict body (e.g.
# `int(config.get("save_steps", 0) or 0)`) and miss any future
# unconditional addition of the gated fields later in the dict.
unconditional_block_start = source.find("mlx_config_kwargs = dict(")
self.assertNotEqual(unconditional_block_start, -1)
depth = 0
i = unconditional_block_start + len("mlx_config_kwargs = dict")
end = i
while i < len(source):
ch = source[i]
if ch == "(":
depth += 1
elif ch == ")":
depth -= 1
if depth == 0:
end = i + 1
break
i += 1
unconditional = source[unconditional_block_start:end]
self.assertNotIn("cast_norm_output_to_input_dtype", unconditional)
self.assertNotIn("dataset_order", unconditional)
self.assertNotIn("max_grad_leaf_norm", unconditional)
self.assertNotIn("append_eos", unconditional)
def test_training_route_forwards_embedding_learning_rate(self):
training_route = _load_route_module(

93
studio/backend/tests/test_training_resume.py Normal file
View file

@ -0,0 +1,93 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression tests for resumable training run eligibility."""
import importlib.util
import json
from pathlib import Path
_BACKEND = Path(__file__).resolve().parents[1]
def _load_resume_module():
spec = importlib.util.spec_from_file_location(
"training_resume_under_test",
_BACKEND / "core" / "training" / "resume.py",
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
resume = _load_resume_module()
def _stopped_run(**overrides):
run = {
"status": "stopped",
"final_step": 5,
"total_steps": 10,
"output_dir": "/tmp/unsloth-output",
"resumed_later": False,
"config_json": json.dumps({"hf_dataset": "org/dataset"}),
}
run.update(overrides)
return run
def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
assert resume.can_resume_run(_stopped_run()) is True
def test_can_resume_run_rejects_s3_dataset_source(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
run = _stopped_run(
config_json = json.dumps(
{
"dataset_source": "s3",
"s3_dataset": {
"bucket": "training-data",
"prefix": "datasets/",
"region": "us-east-1",
"use_iam_role": True,
},
}
)
)
assert resume.can_resume_run(run) is False
def test_can_resume_run_rejects_s3_metadata_marker(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
run = _stopped_run(config_json = json.dumps({"s3_dataset": {"bucket": "training-data"}}))
assert resume.can_resume_run(run) is False
def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
config_json = json.dumps({"dataset_source": "s3", "s3_dataset": {"bucket": "training-data"}})
studio_db.create_run(
id = "run-s3",
model_name = "unsloth/test-model",
dataset_name = "s3://training-data",
config_json = config_json,
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
result = studio_db.list_runs()
assert result["runs"][0]["config_json"] == config_json

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

@ -30,6 +30,7 @@ _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
from utils.models.model_config import (
ModelConfig,
is_vision_model,
_is_vision_model_uncached,
_vision_detection_cache,
@ -120,6 +121,99 @@ class TestVisionCacheSubprocessPath:
mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None)
# ---------------------------------------------------------------------------
# Local GGUF capability path
# ---------------------------------------------------------------------------
class TestLocalGgufVisionDetection:
@patch(
"utils.models.model_config._is_vision_model_subprocess",
side_effect = AssertionError("GGUF must not use Transformers vision detection"),
)
def test_qwen36_gguf_with_mmproj_skips_transformers(self, mock_subprocess, tmp_path):
model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf"
model.write_bytes(b"")
(tmp_path / "mmproj-F32.gguf").write_bytes(b"")
assert is_vision_model(str(model)) is True
mock_subprocess.assert_not_called()
@patch(
"utils.models.model_config._is_vision_model_subprocess",
side_effect = AssertionError("GGUF must not use Transformers vision detection"),
)
def test_direct_gguf_in_variant_subdir_finds_snapshot_mmproj(self, mock_subprocess, tmp_path):
variant_dir = tmp_path / "BF16"
variant_dir.mkdir()
model = variant_dir / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf"
model.write_bytes(b"")
(tmp_path / "mmproj-F32.gguf").write_bytes(b"")
assert is_vision_model(str(model)) is True
mock_subprocess.assert_not_called()
@patch(
"utils.models.model_config._is_vision_model_subprocess",
side_effect = AssertionError("GGUF must not use Transformers vision detection"),
)
def test_qwen36_gguf_without_mmproj_skips_transformers(self, mock_subprocess, tmp_path):
model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf"
model.write_bytes(b"")
assert is_vision_model(str(model)) is False
mock_subprocess.assert_not_called()
def test_local_gguf_check_observes_mmproj_added_later(self, tmp_path):
model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf"
model.write_bytes(b"")
assert is_vision_model(str(model)) is False
(tmp_path / "mmproj-F32.gguf").write_bytes(b"")
assert is_vision_model(str(model)) is True
@patch(
"utils.models.model_config._is_vision_model_subprocess",
side_effect = AssertionError("GGUF must not use Transformers vision detection"),
)
def test_ui_selection_returns_local_gguf_config(self, mock_subprocess, tmp_path):
model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf"
model.write_bytes(b"")
mmproj = tmp_path / "mmproj-F32.gguf"
mmproj.write_bytes(b"")
config = ModelConfig.from_ui_selection(str(model), None)
assert config is not None
assert config.is_gguf is True
assert config.is_vision is True
assert config.gguf_mmproj_file == str(mmproj.resolve())
mock_subprocess.assert_not_called()
@patch(
"utils.models.model_config._is_vision_model_subprocess",
side_effect = AssertionError("GGUF must not use Transformers vision detection"),
)
def test_ui_selection_direct_gguf_in_variant_subdir_keeps_mmproj(
self, mock_subprocess, tmp_path
):
variant_dir = tmp_path / "BF16"
variant_dir.mkdir()
model = variant_dir / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf"
model.write_bytes(b"")
mmproj = tmp_path / "mmproj-F32.gguf"
mmproj.write_bytes(b"")
config = ModelConfig.from_ui_selection(str(model), None)
assert config is not None
assert config.is_gguf is True
assert config.is_vision is True
assert config.gguf_mmproj_file == str(mmproj.resolve())
mock_subprocess.assert_not_called()
# ---------------------------------------------------------------------------
# Exception handling — cache the False fallback

432
studio/backend/utils/hardware/apple.py Normal file
View file

@ -0,0 +1,432 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Apple Silicon GPU temperature and power -- no sudo required.
Mirrors macmon's approach (https://github.com/vladkens/macmon):
* Temperature: average of the AppleSMC "Tg*" float keys (available since
macOS 14; on older systems the keys are absent and this returns None).
* Power: IOReport "Energy Model" group, "GPU Energy" channel. Each poll
diffs the energy counter against the previous poll's sample, so the
result is the average wattage over the polling window. The first poll
only sets the baseline and returns None.
Public API (never raises; returns None when sensors are unavailable):
read_gpu_temperature_c()
read_gpu_power_w()
"""
import ctypes
import struct
import time
from typing import Iterable, Optional
from loggers import get_logger
logger = get_logger(__name__)
_IOKIT_PATH = "/System/Library/Frameworks/IOKit.framework/IOKit"
_CF_PATH = "/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation"
_IOREPORT_PATH = "/usr/lib/libIOReport.dylib"
# AppleSMC user-client protocol (same constants as macmon / SMCKit).
_SMC_SELECTOR_HANDLE_EVENT = 2
_SMC_CMD_READ_BYTES = 5
_SMC_CMD_KEY_AT_INDEX = 8
_SMC_CMD_KEY_INFO = 9
_MAX_VALID_TEMP_C = 150.0
_CF_STRING_ENCODING_UTF8 = 0x08000100
_ENERGY_UNIT_DIVISORS = {"mJ": 1e3, "uJ": 1e6, "nJ": 1e9}
# ========== Pure helpers ==========
def _fourcc(key: str) -> int:
"""Encode a 4-char SMC key/type name as a big-endian integer."""
return int.from_bytes(key.encode("ascii"), "big")
def _fourcc_str(value: int) -> str:
return value.to_bytes(4, "big").decode("ascii", errors = "replace")
def _watts(energy: int, unit: str, elapsed_s: float) -> Optional[float]:
"""Convert an IOReport energy counter delta into average watts."""
divisor = _ENERGY_UNIT_DIVISORS.get(unit.strip())
if divisor is None or elapsed_s <= 0:
return None
return energy / divisor / elapsed_s
def _average_valid_temps(values: Iterable[float]) -> Optional[float]:
valid = [v for v in values if 0.0 < v <= _MAX_VALID_TEMP_C]
if not valid:
return None
return round(sum(valid) / len(valid), 1)
def _is_gpu_energy_channel(name: str) -> bool:
# Exact "GPU Energy" plus "DIE_N_GPU Energy" on Ultra chips; the separate
# "GPU SRAM*" channels are not GPU core power.
return name.endswith("GPU Energy") and "SRAM" not in name
# ========== AppleSMC structs (layout must match the kernel exactly) ==========
class _SMCKeyDataVers(ctypes.Structure):
_fields_ = [
("major", ctypes.c_uint8),
("minor", ctypes.c_uint8),
("build", ctypes.c_uint8),
("reserved", ctypes.c_uint8),
("release", ctypes.c_uint16),
]
class _SMCPLimitData(ctypes.Structure):
_fields_ = [
("version", ctypes.c_uint16),
("length", ctypes.c_uint16),
("cpu_p_limit", ctypes.c_uint32),
("gpu_p_limit", ctypes.c_uint32),
("mem_p_limit", ctypes.c_uint32),
]
class _SMCKeyInfo(ctypes.Structure):
_fields_ = [
("data_size", ctypes.c_uint32),
("data_type", ctypes.c_uint32),
("data_attributes", ctypes.c_uint8),
]
class _SMCKeyData(ctypes.Structure):
_fields_ = [
("key", ctypes.c_uint32),
("vers", _SMCKeyDataVers),
("p_limit_data", _SMCPLimitData),
("key_info", _SMCKeyInfo),
("result", ctypes.c_uint8),
("status", ctypes.c_uint8),
("data8", ctypes.c_uint8),
("data32", ctypes.c_uint32),
("bytes", ctypes.c_uint8 * 32),
]
# ========== Library loaders ==========
def _load_iokit() -> ctypes.CDLL:
iokit = ctypes.CDLL(_IOKIT_PATH)
iokit.IOServiceMatching.restype = ctypes.c_void_p
iokit.IOServiceMatching.argtypes = [ctypes.c_char_p]
iokit.IOServiceGetMatchingServices.argtypes = [
ctypes.c_uint32,
ctypes.c_void_p,
ctypes.POINTER(ctypes.c_uint32),
]
iokit.IOIteratorNext.restype = ctypes.c_uint32
iokit.IOIteratorNext.argtypes = [ctypes.c_uint32]
iokit.IORegistryEntryGetName.argtypes = [ctypes.c_uint32, ctypes.c_char_p]
iokit.IOServiceOpen.argtypes = [
ctypes.c_uint32,
ctypes.c_uint32,
ctypes.c_uint32,
ctypes.POINTER(ctypes.c_uint32),
]
iokit.IOObjectRelease.argtypes = [ctypes.c_uint32]
iokit.IOConnectCallStructMethod.argtypes = [
ctypes.c_uint32,
ctypes.c_uint32,
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_void_p,
ctypes.POINTER(ctypes.c_size_t),
]
return iokit
def _load_cf() -> ctypes.CDLL:
cf = ctypes.CDLL(_CF_PATH)
cf.CFStringCreateWithCString.restype = ctypes.c_void_p
cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint32]
cf.CFStringGetCString.restype = ctypes.c_bool
cf.CFStringGetCString.argtypes = [
ctypes.c_void_p,
ctypes.c_char_p,
ctypes.c_long,
ctypes.c_uint32,
]
cf.CFRelease.argtypes = [ctypes.c_void_p]
cf.CFDictionaryGetValue.restype = ctypes.c_void_p
cf.CFDictionaryGetValue.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
cf.CFArrayGetCount.restype = ctypes.c_long
cf.CFArrayGetCount.argtypes = [ctypes.c_void_p]
cf.CFArrayGetValueAtIndex.restype = ctypes.c_void_p
cf.CFArrayGetValueAtIndex.argtypes = [ctypes.c_void_p, ctypes.c_long]
return cf
def _load_ioreport() -> ctypes.CDLL:
ior = ctypes.CDLL(_IOREPORT_PATH)
ior.IOReportCopyChannelsInGroup.restype = ctypes.c_void_p
ior.IOReportCopyChannelsInGroup.argtypes = [
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_uint64,
ctypes.c_uint64,
ctypes.c_uint64,
]
ior.IOReportCreateSubscription.restype = ctypes.c_void_p
ior.IOReportCreateSubscription.argtypes = [
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.POINTER(ctypes.c_void_p),
ctypes.c_uint64,
ctypes.c_void_p,
]
ior.IOReportCreateSamples.restype = ctypes.c_void_p
ior.IOReportCreateSamples.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
ior.IOReportCreateSamplesDelta.restype = ctypes.c_void_p
ior.IOReportCreateSamplesDelta.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
ior.IOReportChannelGetChannelName.restype = ctypes.c_void_p
ior.IOReportChannelGetChannelName.argtypes = [ctypes.c_void_p]
ior.IOReportChannelGetUnitLabel.restype = ctypes.c_void_p
ior.IOReportChannelGetUnitLabel.argtypes = [ctypes.c_void_p]
ior.IOReportSimpleGetIntegerValue.restype = ctypes.c_int64
ior.IOReportSimpleGetIntegerValue.argtypes = [ctypes.c_void_p, ctypes.c_int32]
return ior
def _cfstr(cf: ctypes.CDLL, text: str) -> int:
return cf.CFStringCreateWithCString(None, text.encode("utf-8"), _CF_STRING_ENCODING_UTF8)
def _from_cfstr(cf: ctypes.CDLL, ref: Optional[int]) -> str:
if not ref:
return ""
buf = ctypes.create_string_buffer(128)
if not cf.CFStringGetCString(ref, buf, len(buf), _CF_STRING_ENCODING_UTF8):
return ""
return buf.value.decode("utf-8", errors = "replace").strip()
# ========== SMC connection (GPU temperature) ==========
class _SMCConnection:
"""Connection to AppleSMCKeysEndpoint; discovers "Tg*" GPU temp keys once."""
def __init__(self):
self._iokit = _load_iokit()
self._conn = self._open()
self._key_info_cache: dict[int, _SMCKeyInfo] = {}
self.gpu_keys = [
key
for key in self._all_keys()
if key.startswith("Tg") and self.read_float(key) is not None
]
def _open(self) -> int:
iterator = ctypes.c_uint32(0)
matching = self._iokit.IOServiceMatching(b"AppleSMC")
if self._iokit.IOServiceGetMatchingServices(0, matching, ctypes.byref(iterator)) != 0:
raise OSError("AppleSMC service not found")
try:
conn = self._open_keys_endpoint(iterator.value)
finally:
self._iokit.IOObjectRelease(iterator.value)
if conn is None:
raise OSError("AppleSMCKeysEndpoint not found")
return conn
def _open_keys_endpoint(self, iterator: int) -> Optional[int]:
task = ctypes.CDLL(None).mach_task_self()
while device := self._iokit.IOIteratorNext(iterator):
name = ctypes.create_string_buffer(128)
self._iokit.IORegistryEntryGetName(device, name)
if name.value != b"AppleSMCKeysEndpoint":
self._iokit.IOObjectRelease(device)
continue
conn = ctypes.c_uint32(0)
status = self._iokit.IOServiceOpen(device, task, 0, ctypes.byref(conn))
self._iokit.IOObjectRelease(device)
if status != 0:
raise OSError(f"IOServiceOpen(AppleSMCKeysEndpoint) failed: {status}")
return conn.value
return None
def _call(self, ival: _SMCKeyData) -> _SMCKeyData:
oval = _SMCKeyData()
olen = ctypes.c_size_t(ctypes.sizeof(_SMCKeyData))
status = self._iokit.IOConnectCallStructMethod(
self._conn,
_SMC_SELECTOR_HANDLE_EVENT,
ctypes.byref(ival),
ctypes.sizeof(_SMCKeyData),
ctypes.byref(oval),
ctypes.byref(olen),
)
if status != 0:
raise OSError(f"IOConnectCallStructMethod failed: {status}")
if oval.result != 0:
raise OSError(f"SMC result code: {oval.result}")
return oval
def _read_key_info(self, key_id: int) -> _SMCKeyInfo:
cached = self._key_info_cache.get(key_id)
if cached is not None:
return cached
oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_KEY_INFO))
self._key_info_cache[key_id] = oval.key_info
return oval.key_info
def _read_bytes(self, key: str) -> Optional[bytes]:
try:
key_id = _fourcc(key)
info = self._read_key_info(key_id)
oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_READ_BYTES, key_info = info))
return bytes(oval.bytes[: info.data_size])
except OSError:
return None
def read_float(self, key: str) -> Optional[float]:
try:
info = self._read_key_info(_fourcc(key))
except OSError:
return None
if info.data_size != 4 or info.data_type != _fourcc("flt "):
return None
data = self._read_bytes(key)
if data is None or len(data) != 4:
return None
return struct.unpack("<f", data)[0]
def _key_name_at(self, index: int) -> Optional[str]:
try:
oval = self._call(_SMCKeyData(data8 = _SMC_CMD_KEY_AT_INDEX, data32 = index))
return oval.key.to_bytes(4, "big").decode("ascii")
except (OSError, UnicodeDecodeError):
return None
def _all_keys(self) -> list[str]:
count_bytes = self._read_bytes("#KEY")
if count_bytes is None or len(count_bytes) != 4:
return []
count = int.from_bytes(count_bytes, "big")
names = (self._key_name_at(i) for i in range(count))
return [name for name in names if name is not None]
def gpu_temperature_c(self) -> Optional[float]:
readings = (self.read_float(key) for key in self.gpu_keys)
return _average_valid_temps(value for value in readings if value is not None)
# ========== IOReport subscription (GPU power) ==========
class _IOReportEnergy:
"""Persistent subscription to the "Energy Model" group for GPU wattage."""
def __init__(self):
self._cf = _load_cf()
self._ior = _load_ioreport()
self._channels = self._ior.IOReportCopyChannelsInGroup(
_cfstr(self._cf, "Energy Model"), None, 0, 0, 0
)
if not self._channels:
raise OSError("IOReport 'Energy Model' channel group unavailable")
subscribed = ctypes.c_void_p()
self._sub = self._ior.IOReportCreateSubscription(
None, self._channels, ctypes.byref(subscribed), 0, None
)
if not self._sub:
raise OSError("IOReportCreateSubscription failed")
# Sample with the channels IOReport subscribes us to, not the requested
# group (matches macmon); fall back if the OS leaves it unset.
self._sample_channels = subscribed if subscribed else self._channels
self._channels_key = _cfstr(self._cf, "IOReportChannels")
self._prev: Optional[tuple[int, float]] = None # (sample ref, monotonic s)
def gpu_power_w(self) -> Optional[float]:
sample = self._ior.IOReportCreateSamples(self._sub, self._sample_channels, None)
if not sample:
return None
now = time.monotonic()
prev, self._prev = self._prev, (sample, now)
if prev is None:
return None
prev_sample, prev_time = prev
delta = self._ior.IOReportCreateSamplesDelta(prev_sample, sample, None)
self._cf.CFRelease(prev_sample)
if not delta:
return None
try:
return self._gpu_watts_from_delta(delta, now - prev_time)
finally:
self._cf.CFRelease(delta)
def _gpu_watts_from_delta(self, delta: int, elapsed_s: float) -> Optional[float]:
items = self._cf.CFDictionaryGetValue(delta, self._channels_key)
if not items:
return None
total: Optional[float] = None
for i in range(self._cf.CFArrayGetCount(items)):
item = self._cf.CFArrayGetValueAtIndex(items, i)
name = _from_cfstr(self._cf, self._ior.IOReportChannelGetChannelName(item))
if not _is_gpu_energy_channel(name):
continue
unit = _from_cfstr(self._cf, self._ior.IOReportChannelGetUnitLabel(item))
energy = self._ior.IOReportSimpleGetIntegerValue(item, 0)
watts = _watts(energy, unit, elapsed_s)
if watts is not None:
total = (total or 0.0) + watts
if total is None or total < 0: # negative = counter reset; show -- not a bogus draw
return None
return round(total, 1)
# ========== Public API (module singletons, failure-latched) ==========
_smc: Optional[_SMCConnection] = None
_smc_failed = False
_energy: Optional[_IOReportEnergy] = None
_energy_failed = False
def read_gpu_temperature_c() -> Optional[float]:
"""Average Apple GPU die temperature in degrees C, or None if unavailable."""
global _smc, _smc_failed
if _smc_failed:
return None
try:
if _smc is None:
_smc = _SMCConnection()
return _smc.gpu_temperature_c()
except Exception as e:
_smc_failed = True
logger.warning("Apple SMC GPU temperature unavailable: %s", e)
return None
def read_gpu_power_w() -> Optional[float]:
"""Average GPU power in watts since the previous call, or None.
The first call establishes the baseline sample and returns None.
"""
global _energy, _energy_failed
if _energy_failed:
return None
try:
if _energy is None:
_energy = _IOReportEnergy()
return _energy.gpu_power_w()
except Exception as e:
_energy_failed = True
logger.warning("Apple IOReport GPU power unavailable: %s", e)
return None

6
studio/backend/utils/hardware/hardware.py
View file

@ -714,17 +714,19 @@ def get_gpu_utilization() -> Dict[str, Any]:
except Exception:
pass
from . import apple
return {
"available": True,
"backend": device.value,
"gpu_utilization_pct": agx.get("utilization_pct") if agx else None,
"temperature_c": None,
"temperature_c": apple.read_gpu_temperature_c(),
"vram_used_gb": round(vram_used_gb, 2),
"vram_total_gb": round(total_gb, 2),
"vram_utilization_pct": (
round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None
),
"power_draw_w": None,
"power_draw_w": apple.read_gpu_power_w(),
"power_limit_w": None,
"power_utilization_pct": None,
}

78
studio/backend/utils/host_policy.py Normal file
View file

@ -0,0 +1,78 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Bind-host trust policy for the Studio backend.
Stdlib only -- safe to import without the rest of the backend.
`is_external_host` mirrors the CLI's `unsloth_cli/_tool_policy.py`: a loopback
bind is the user's own machine, any other address is network-reachable. The
logic is duplicated rather than shared because the backend is self-contained
(see run.py: "can be moved to any directory") and runs from a venv that may not
have `unsloth_cli` on sys.path. Keep the two in sync.
"""
from __future__ import annotations
import os
# Loopback aliases; any other bind address is treated as network-reachable. Only
# the exact aliases the rest of the stack assumes for loopback (health checks,
# banner URLs, run.py all hard-code 127.0.0.1), so other 127.0.0.0/8 addresses
# are deliberately left out -- they are not supported launch hosts.
_LOOPBACK_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"})
# Whether a loopback launch in THIS process auto-enabled the gate. run_server
# normally runs once per process, but if it is reused with a different host
# (embedders, tests) a stale loopback default must not carry into a later
# public bind, so we only ever take back a value we set ourselves.
_auto_enabled = False
def is_external_host(host: str) -> bool:
"""True when `host` is reachable from beyond loopback."""
return host.lower() not in _LOOPBACK_HOSTS
def apply_stdio_mcp_loopback_default(host: str, *, is_colab: bool = False) -> None:
"""Default stdio MCP servers on when bound to loopback.
A loopback bind is the user's own machine -- the same trust boundary the
Tauri desktop app relies on (see main.py, which also binds 127.0.0.1 and
setdefaults this var). Colab is excluded: even its loopback is a hosted VM
reachable through Colab's proxy, so it stays off unless opted in. An explicit
operator value wins: a pre-set `UNSLOTH_STUDIO_ALLOW_STDIO_MCP=0`
force-disables and `=1` opts in, including on a network bind. We only ever
set or clear a default we applied ourselves, so reusing run_server with a
public host after a loopback one does not leave the gate on.
"""
global _auto_enabled
current = os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP")
# If our prior auto-default was changed out from under us (in-process reuse),
# relinquish ownership: an explicit =0 is then honored below as a sticky
# force-disable, while a cleared var falls back to the host default like a
# fresh process.
if _auto_enabled and current != "1":
_auto_enabled = False
# An explicit operator value is one we did not set; never touch it.
if current is not None and not _auto_enabled:
return
if is_colab or is_external_host(host):
if _auto_enabled:
os.environ.pop("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", None)
_auto_enabled = False
else:
os.environ["UNSLOTH_STUDIO_ALLOW_STDIO_MCP"] = "1"
_auto_enabled = True
def loopback_default_active() -> bool:
"""True when stdio MCP is on only because a loopback bind auto-enabled it,
rather than an explicit operator opt-in. Lets the gate tell the two apart."""
return _auto_enabled
def _reset_loopback_default_state() -> None:
"""Test hook: forget any auto-enable applied earlier in this process."""
global _auto_enabled
_auto_enabled = False

19
studio/backend/utils/llama_cpp_freshness.py
View file

@ -301,7 +301,22 @@ def format_stale_warning(info: dict) -> str:
)
def reset_caches() -> None:
"""Test-only: drop all in-memory caches."""
def reset_caches(*, drop_disk: bool = False) -> None:
"""Drop the in-memory freshness caches. The no-arg form is test-only.
With ``drop_disk = True`` also delete the on-disk 24h release cache. Used by
the post-install/update path: in-memory clearing alone leaves the stale
same-base value on disk, so if the post-install GitHub refresh can't reach
the network, ``latest_published_release`` would replay that stale disk value
(see its last-good fallback) and the banner could linger. Dropping the disk
cache makes latest read as None in that offline case, so the banner fails
open (off) instead of pointing at the just-replaced build."""
_marker_cache.clear()
_release_memo.clear()
if drop_disk:
import shutil
# _cache_dir() is a dedicated freshness-only subdir; it is re-created on
# the next _save_disk_cache. ignore_errors so a missing/locked dir is a
# no-op rather than breaking an otherwise successful install.
shutil.rmtree(_cache_dir(), ignore_errors = True)

56
studio/backend/utils/llama_cpp_update.py
View file

@ -37,6 +37,7 @@ from utils.llama_cpp_freshness import (
_INSTALL_MARKER_NAME,
check_prebuilt_freshness,
latest_published_release,
parse_base_build,
read_install_marker,
reset_caches,
)
@ -220,23 +221,42 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
res = _resolve_prebuilt_for_host(force_refresh = force_refresh)
if not res or not res.get("prebuilt_available"):
return None
# llama_tag is the upstream build (bNNNN, what --version reports); release_tag
# can be a fork wrapper tag, so compare/display against llama_tag.
latest = res.get("llama_tag") or res.get("release_tag")
if not latest:
# llama_tag is the upstream base (bNNNN, what --version reports); release_tag
# is the full tag, either a same-base mix (bNNNN-mix-<sha>) or a fork wrapper
# (e.g. v1.0). Compare the numeric base against llama_tag.
base_tag = res.get("llama_tag") or res.get("release_tag")
release_tag = res.get("release_tag")
if not base_tag:
return None
# No resolvable install root (e.g. a pinned LLAMA_SERVER_PATH we cannot
# manage) means an apply would not take effect, so do not offer.
if _llama_install_root(binary) is None:
return None
installed_build = _installed_build_number(binary)
m = re.search(r"(\d+)", latest)
latest_build = int(m.group(1)) if m else None
# Suppress only when the source build is reliably newer/equal; unknown
# version (the involuntary source-build case) is treated as behind.
update_available = (
installed_build is None or latest_build is None or installed_build < latest_build
latest_build = parse_base_build(base_tag)
# A same-base mix adds patches the bare base lacks, so it is newer even at an
# unchanged build number (the marker path's is_behind already does this). The
# bNNNN anchor keeps a fork wrapper tag from being read as a mix.
latest_is_mix = (
isinstance(release_tag, str)
and latest_build is not None
and parse_base_build(release_tag) == latest_build
and release_tag.strip() != f"b{latest_build}"
)
if installed_build is None or latest_build is None:
# Unknown installed/latest version (the involuntary source-build case):
# treat as behind so we still offer the prebuilt.
update_available = True
elif installed_build < latest_build:
update_available = True
elif installed_build == latest_build:
# Same upstream base: offer the extra-patch mix, never a bare rebuild.
update_available = latest_is_mix
else:
# Source build newer than the latest prebuilt: downgrade guard.
update_available = False
# Display the mix tag when that's what makes it newer; otherwise the base.
latest = release_tag if latest_is_mix else base_tag
with _job_lock:
job = dict(_job)
return {
@ -310,8 +330,9 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
def _rocm_install_args(asset: Optional[str]) -> list[str]:
"""Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh.
The installer probe can miss the gfx arch on amd-smi-only hosts; lemonade
bundles carry the family in the name (rocm-gfx110X), fork bundles only rocm/hip."""
The installer probe can miss the gfx arch on amd-smi-only hosts; per-gfx
ROCm bundles carry the family in the name (rocm-gfx110X), version-tagged
bundles only rocm/hip."""
if not asset:
return []
low = asset.lower()
@ -406,10 +427,13 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
tail = "".join(tail_lines).strip()[-1500:]
raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}")
# New UNSLOTH_PREBUILT_INFO.json is on disk; drop in-memory caches and
# re-prime the 24h disk freshness cache with the true newest, so the
# banner can't linger on a stale same-base value after the swap.
reset_caches()
# New UNSLOTH_PREBUILT_INFO.json is on disk; drop the in-memory AND the
# on-disk freshness caches, then re-prime the 24h disk cache with the
# true newest, so the banner can't linger on a stale same-base value
# after the swap. drop_disk matters when the refresh below can't reach
# GitHub: without it, latest_published_release would replay the stale
# disk value; with it, latest reads as None and the banner fails open.
reset_caches(drop_disk = True)
try:
latest_published_release(repo, force_refresh = True)
except Exception as exc: # pragma: no cover - network defensive

92
studio/backend/utils/models/model_config.py
View file

@ -721,6 +721,25 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
model_name: Model identifier (HF repo or local path)
hf_token: Optional HF token for gated/private models
"""
# Local GGUF models are served by llama-server. Their multimodal
# capability comes from a companion mmproj, not a Transformers config.
# Do not cache this lookup: a projector may be added beside an existing
# weight file after it was first inspected.
if is_local_path(model_name):
local_path = normalize_path(model_name)
gguf_file = detect_gguf_model(local_path)
if gguf_file:
companion_root = _local_gguf_companion_search_root(local_path, gguf_file)
mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root)
is_vision = mmproj_file is not None
logger.debug(
"Local GGUF vision check for '%s': mmproj=%s, is_vision=%s",
gguf_file,
mmproj_file,
is_vision,
)
return is_vision
# Normalize model name so different casings of the same repo share a key
try:
if is_local_path(model_name):
@ -1324,6 +1343,7 @@ def _extract_quant_label(filename: str) -> str:
"model-UD-IQ1_S.gguf" → "UD-IQ1_S"
"model-UD-TQ1_0.gguf" → "UD-TQ1_0"
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE"
"Qwen3.6-IQ4_XS-3.53bpw.gguf" → "IQ4_XS-3.53bpw"
"""
import re
@ -1339,6 +1359,10 @@ def _extract_quant_label(filename: str) -> str:
r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1
r"|Q[0-9]+_K" # Short K-quant: Q6_K
r"|BF16|F16|F32)" # Full precision
# Optional bits-per-weight modifier so repos that ship multiple
# files at the same base quant (e.g. byteshape's IQ4_XS at 3.53,
# 3.97, 4.19 bpw) don't collapse into a single merged variant.
r"(-[0-9]+(?:\.[0-9]+)?bpw)?"
)
match = re.search(quant_re, stem, re.IGNORECASE)
# Subdir layouts like ``BF16/foo.gguf`` keep the quant in the directory,
@ -1353,11 +1377,41 @@ def _extract_quant_label(filename: str) -> str:
break
if match:
prefix = match.group(1) or ""
return f"{prefix}{match.group(2)}"
bpw = match.group(3) or ""
return f"{prefix}{match.group(2)}{bpw}"
# Fallback: last hyphen-separated segment
return stem.split("-")[-1]
def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str:
"""Directory to scan upward from for local GGUF companion files."""
import re
selected = Path(selected_path)
gguf_path = Path(gguf_file)
if selected.suffix.lower() != ".gguf":
return selected_path
gguf_dir = gguf_path.parent
if not gguf_dir.name:
return str(gguf_dir)
quant_dir_re = (
r"(UD-)?("
r"MXFP[0-9]+(?:_[A-Z0-9]+)*"
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
r"|TQ[0-9]+_[0-9]+"
r"|Q[0-9]+_K_[A-Z]+"
r"|Q[0-9]+_[0-9]+"
r"|Q[0-9]+_K"
r"|BF16|F16|F32"
r")"
)
if re.fullmatch(quant_dir_re, gguf_dir.name, re.IGNORECASE):
return str(gguf_dir.parent)
return str(gguf_dir)
def _iter_hf_cache_snapshots(repo_id: str):
"""Yield HF cache snapshot dirs for *repo_id*, newest first.
@ -1371,12 +1425,11 @@ def _iter_hf_cache_snapshots(repo_id: str):
return
cache_dir = Path(hf_constants.HF_HUB_CACHE)
if not cache_dir.is_dir():
return
target = f"models--{repo_id.replace('/', '--')}".lower()
repo_dir: Optional[Path] = None
try:
if not cache_dir.is_dir():
return
for entry in cache_dir.iterdir():
if entry.is_dir() and entry.name.lower() == target:
repo_dir = entry
@ -1387,10 +1440,9 @@ def _iter_hf_cache_snapshots(repo_id: str):
return
snapshots = repo_dir / "snapshots"
if not snapshots.is_dir():
return
try:
if not snapshots.is_dir():
return
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()]
except OSError:
return
@ -2277,10 +2329,10 @@ class ModelConfig:
except Exception as e:
logger.debug(f"Could not read export metadata: {e}")
# Pass search_root=path so detect_mmproj_file walks up to the
# snapshot root: the weight may sit in a quant subdir while
# mmproj-*.gguf lives at the root.
mmproj_file = detect_mmproj_file(gguf_file, search_root = path)
# Direct file selections may point into a quant subdir while
# mmproj-*.gguf lives at the snapshot root.
companion_root = _local_gguf_companion_search_root(path, gguf_file)
mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root)
if mmproj_file:
gguf_is_vision = True
logger.info(f"Detected mmproj for vision: {mmproj_file}")
@ -2288,7 +2340,7 @@ class ModelConfig:
logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
# Separate MTP drafter sibling (Gemma 4), mirroring mmproj.
mtp_file = detect_mtp_file(gguf_file, search_root = path)
mtp_file = detect_mtp_file(gguf_file, search_root = companion_root)
if mtp_file:
logger.info(f"Detected MTP drafter: {mtp_file}")
@ -2309,10 +2361,13 @@ class ModelConfig:
# Does the HF repo contain GGUF files?
gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token)
if gguf_filename:
# Preflight: verify llama-server binary exists before a multi-GB download
# Preflight: verify llama-server binary exists before a multi-GB
# download. include_denied: a transiently locked binary still
# exists (the lock clears long before the download finishes; the
# load itself reports a still-locked binary distinctly).
from core.inference.llama_cpp import LlamaCppBackend
if not LlamaCppBackend._find_llama_server_binary():
if not LlamaCppBackend._find_llama_server_binary(include_denied = True):
raise RuntimeError(
"llama-server binary not found — cannot load GGUF models. "
"Run setup.sh to build it, or set LLAMA_SERVER_PATH."
@ -2478,6 +2533,15 @@ class ModelConfig:
identifier = resolved_identifier
path = resolved_identifier
# Keep existing local GGUF selections on the llama-server path. This
# constructor is still used by older inference helpers and must not
# describe a .gguf weight file as loadable by FastVisionModel.
if is_local and not is_lora and detect_gguf_model(path):
gguf_config = cls.from_identifier(path, hf_token = hf_token)
if gguf_config is not None:
gguf_config.display_name = display_name
return gguf_config
# --- Base Model and Vision Detection ---
base_model = None
is_vision = False

6
studio/backend/utils/paths/__init__.py
View file

@ -41,15 +41,17 @@ from .storage_roots import (
ensure_dir,
ensure_studio_directories,
resolve_under_root,
default_run_dir_name,
resolve_output_dir,
resolve_export_dir,
resolve_export_write_dir,
resolve_tensorboard_dir,
resolve_dataset_path,
)
# Re-export shim: mark project-path helpers as used so the import-hoist
# safety net does not flag them as unused.
_REEXPORTED = (documents_root, project_workspaces_root)
_REEXPORTED = (documents_root, project_workspaces_root, resolve_export_write_dir)
__all__ = [
"normalize_path",
@ -87,8 +89,10 @@ __all__ = [
"ensure_dir",
"ensure_studio_directories",
"resolve_under_root",
"default_run_dir_name",
"resolve_output_dir",
"resolve_export_dir",
"resolve_export_write_dir",
"resolve_tensorboard_dir",
"resolve_dataset_path",
]

74
studio/backend/utils/paths/storage_roots.py
View file

@ -5,8 +5,9 @@ from __future__ import annotations
import json
import os
import re
import sys
from pathlib import Path
from pathlib import Path, PurePosixPath, PureWindowsPath
import tempfile
@ -317,6 +318,26 @@ def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = (
return Path(*parts) if parts else Path()
def _has_parent_segment(raw: str, path: Path) -> bool:
"""Return true when a user path contains a parent-directory segment.
On POSIX, ``Path("E:\\foo\\..\\bar")`` treats backslashes as normal
characters, so check both the host parser and Windows-style parsing.
"""
if ".." in path.parts:
return True
if ".." in PureWindowsPath(raw).parts:
return True
return ".." in raw.replace("\\", "/").split("/")
def _is_absolute_user_path(path: Path) -> bool:
expanded = str(path)
if os.name == "nt":
return path.is_absolute() and PureWindowsPath(expanded).is_absolute()
return path.is_absolute() and PurePosixPath(expanded).is_absolute()
def _assert_contained(resolved: Path, root: Path) -> None:
"""Raise ValueError if ``resolved`` realpaths outside ``root``."""
try:
@ -351,10 +372,10 @@ def resolve_under_root(
raise ValueError("path may not contain null bytes")
path = Path(raw).expanduser()
if ".." in path.parts:
if _has_parent_segment(raw, path):
raise ValueError(f"path may not contain '..' segments: {raw!r}")
if path.is_absolute():
if _is_absolute_user_path(path):
_assert_contained(path, root)
return path
@ -364,6 +385,23 @@ def resolve_under_root(
return candidate
def default_run_dir_name(model_name: str) -> str:
# Folder-safe run name for an auto-created output dir. Repo ids keep their
# namespace (org/model -> org_model); local paths (incl. G:\dir\model)
# collapse to their final component so an absolute source can't escape
# outputs_root. Length-capped to stay under the filesystem name limit.
raw = str(model_name or "").strip()
is_path = (
"\\" in raw
or raw.startswith(("/", "~", "."))
or os.path.isabs(raw)
or (len(raw) >= 2 and raw[1] == ":")
)
base = PureWindowsPath(raw).name if is_path else raw.replace("/", "_")
base = re.sub(r"[^A-Za-z0-9._-]+", "_", base)[:200].strip("._-")
return base or "model"
def resolve_output_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
@ -373,6 +411,36 @@ def resolve_output_dir(path_value: str | None = None) -> Path:
def resolve_export_dir(path_value: str | None = None) -> Path:
"""Resolve an export directory — contained under exports_root().
Used by scan/read endpoints. Use :func:`resolve_export_write_dir`
for the export write path where absolute paths are accepted.
"""
return resolve_under_root(
path_value,
root = exports_root(),
strip_prefixes = ("exports",),
)
def resolve_export_write_dir(path_value: str | None = None) -> Path:
"""Resolve an export save directory — accepts absolute paths.
Unlike :func:`resolve_export_dir`, this function passes absolute
paths through as-is so users can target a different drive when
their Studio install lives on a constrained system volume
(see :gh-issue:`6082`). Used only by the export write path.
"""
if not path_value or not str(path_value).strip():
return exports_root()
raw = str(path_value).strip()
if "\x00" in raw:
raise ValueError("path may not contain null bytes")
path = Path(raw).expanduser()
if _has_parent_segment(raw, path):
raise ValueError(f"path may not contain '..' segments: {raw!r}")
if _is_absolute_user_path(path):
return path
return resolve_under_root(
path_value,
root = exports_root(),

26
studio/frontend/index.html
View file

@ -1,16 +1,16 @@
<!doctype html>
<!doctype html>
<!-- SPDX-License-Identifier: AGPL-3.0-only -->
<!-- Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -->
<html lang="en">
<head>
<meta charset="UTF-8" />
<link rel="icon" type="image/png" href="/favicon.png" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>Unsloth Studio</title>
</head>
<body>
<div id="root"></div>
<script type="module" src="/src/main.tsx"></script>
</body>
</html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<link rel="icon" type="image/png" href="/favicon.png" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>Unsloth Studio</title>
</head>
<body>
<div id="root"></div>
<script type="module" src="/src/main.tsx"></script>
</body>
</html>

36
studio/frontend/public/hub/profile/logo/meta.svg
View file

@ -1,19 +1,19 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns:cc="http://creativecommons.org/ns#" width="287.56" height="191">
<desc>Logo of Meta Platforms -- Graphic created by Detmar Owen</desc>
<defs>
<linearGradient id="Grad_Logo1" x1="61" y1="117" x2="259" y2="127" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0064e1" offset="0"/>
<stop style="stop-color:#0064e1" offset="0.4"/>
<stop style="stop-color:#0073ee" offset="0.83"/>
<stop style="stop-color:#0082fb" offset="1"/>
</linearGradient>
<linearGradient id="Grad_Logo2" x1="45" y1="139" x2="45" y2="66" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0082fb" offset="0"/>
<stop style="stop-color:#0064e0" offset="1"/>
</linearGradient>
</defs>
<path id="Logo0" style="fill:#0081fb" d="m31.06,125.96c0,10.98 2.41,19.41 5.56,24.51 4.13,6.68 10.29,9.51 16.57,9.51 8.1,0 15.51-2.01 29.79-21.76 11.44-15.83 24.92-38.05 33.99-51.98l15.36-23.6c10.67-16.39 23.02-34.61 37.18-46.96 11.56-10.08 24.03-15.68 36.58-15.68 21.07,0 41.14,12.21 56.5,35.11 16.81,25.08 24.97,56.67 24.97,89.27 0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75l0-31.02c17.63,0 22.03-16.2 22.03-34.74 0-26.42-6.16-55.74-19.73-76.69-9.63-14.86-22.11-23.94-35.84-23.94-14.85,0-26.8,11.2-40.23,31.17-7.14,10.61-14.47,23.54-22.7,38.13l-9.06,16.05c-18.2,32.27-22.81,39.62-31.91,51.75-15.95,21.24-29.57,29.29-47.5,29.29-21.27,0-34.72-9.21-43.05-23.09-6.8-11.31-10.14-26.15-10.14-43.06z"/>
<path id="Logo1" style="fill:url(#Grad_Logo1)" d="m24.49,37.3c14.24-21.95 34.79-37.3 58.36-37.3 13.65,0 27.22,4.04 41.39,15.61 15.5,12.65 32.02,33.48 52.63,67.81l7.39,12.32c17.84,29.72 27.99,45.01 33.93,52.22 7.64,9.26 12.99,12.02 19.94,12.02 17.63,0 22.03-16.2 22.03-34.74l27.4-.86c0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75-12.8,0-24.14-2.78-36.68-14.61-9.64-9.08-20.91-25.21-29.58-39.71l-25.79-43.08c-12.94-21.62-24.81-37.74-31.68-45.04-7.39-7.85-16.89-17.33-32.05-17.33-12.27,0-22.69,8.61-31.41,21.78z"/>
<path id="Logo2" style="fill:url(#Grad_Logo2)" d="m82.35,31.23c-12.27,0-22.69,8.61-31.41,21.78-12.33,18.61-19.88,46.33-19.88,72.95 0,10.98 2.41,19.41 5.56,24.51l-26.48,17.44c-6.8-11.31-10.14-26.15-10.14-43.06 0-30.75 8.44-62.8 24.49-87.55 14.24-21.95 34.79-37.3 58.36-37.3z"/>
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns:cc="http://creativecommons.org/ns#" width="287.56" height="191">
<desc>Logo of Meta Platforms -- Graphic created by Detmar Owen</desc>
<defs>
<linearGradient id="Grad_Logo1" x1="61" y1="117" x2="259" y2="127" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0064e1" offset="0"/>
<stop style="stop-color:#0064e1" offset="0.4"/>
<stop style="stop-color:#0073ee" offset="0.83"/>
<stop style="stop-color:#0082fb" offset="1"/>
</linearGradient>
<linearGradient id="Grad_Logo2" x1="45" y1="139" x2="45" y2="66" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0082fb" offset="0"/>
<stop style="stop-color:#0064e0" offset="1"/>
</linearGradient>
</defs>
<path id="Logo0" style="fill:#0081fb" d="m31.06,125.96c0,10.98 2.41,19.41 5.56,24.51 4.13,6.68 10.29,9.51 16.57,9.51 8.1,0 15.51-2.01 29.79-21.76 11.44-15.83 24.92-38.05 33.99-51.98l15.36-23.6c10.67-16.39 23.02-34.61 37.18-46.96 11.56-10.08 24.03-15.68 36.58-15.68 21.07,0 41.14,12.21 56.5,35.11 16.81,25.08 24.97,56.67 24.97,89.27 0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75l0-31.02c17.63,0 22.03-16.2 22.03-34.74 0-26.42-6.16-55.74-19.73-76.69-9.63-14.86-22.11-23.94-35.84-23.94-14.85,0-26.8,11.2-40.23,31.17-7.14,10.61-14.47,23.54-22.7,38.13l-9.06,16.05c-18.2,32.27-22.81,39.62-31.91,51.75-15.95,21.24-29.57,29.29-47.5,29.29-21.27,0-34.72-9.21-43.05-23.09-6.8-11.31-10.14-26.15-10.14-43.06z"/>
<path id="Logo1" style="fill:url(#Grad_Logo1)" d="m24.49,37.3c14.24-21.95 34.79-37.3 58.36-37.3 13.65,0 27.22,4.04 41.39,15.61 15.5,12.65 32.02,33.48 52.63,67.81l7.39,12.32c17.84,29.72 27.99,45.01 33.93,52.22 7.64,9.26 12.99,12.02 19.94,12.02 17.63,0 22.03-16.2 22.03-34.74l27.4-.86c0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75-12.8,0-24.14-2.78-36.68-14.61-9.64-9.08-20.91-25.21-29.58-39.71l-25.79-43.08c-12.94-21.62-24.81-37.74-31.68-45.04-7.39-7.85-16.89-17.33-32.05-17.33-12.27,0-22.69,8.61-31.41,21.78z"/>
<path id="Logo2" style="fill:url(#Grad_Logo2)" d="m82.35,31.23c-12.27,0-22.69,8.61-31.41,21.78-12.33,18.61-19.88,46.33-19.88,72.95 0,10.98 2.41,19.41 5.56,24.51l-26.48,17.44c-6.8-11.31-10.14-26.15-10.14-43.06 0-30.75 8.44-62.8 24.49-87.55 14.24-21.95 34.79-37.3 58.36-37.3z"/>
</svg>
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 2.3 KiB

After

Width:  |  Height:  |  Size: 2.3 KiB

Before After
Before After

36
studio/frontend/public/provider-logos/misc/meta.svg
View file

@ -1,19 +1,19 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns:cc="http://creativecommons.org/ns#" width="287.56" height="191">
<desc>Logo of Meta Platforms -- Graphic created by Detmar Owen</desc>
<defs>
<linearGradient id="Grad_Logo1" x1="61" y1="117" x2="259" y2="127" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0064e1" offset="0"/>
<stop style="stop-color:#0064e1" offset="0.4"/>
<stop style="stop-color:#0073ee" offset="0.83"/>
<stop style="stop-color:#0082fb" offset="1"/>
</linearGradient>
<linearGradient id="Grad_Logo2" x1="45" y1="139" x2="45" y2="66" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0082fb" offset="0"/>
<stop style="stop-color:#0064e0" offset="1"/>
</linearGradient>
</defs>
<path id="Logo0" style="fill:#0081fb" d="m31.06,125.96c0,10.98 2.41,19.41 5.56,24.51 4.13,6.68 10.29,9.51 16.57,9.51 8.1,0 15.51-2.01 29.79-21.76 11.44-15.83 24.92-38.05 33.99-51.98l15.36-23.6c10.67-16.39 23.02-34.61 37.18-46.96 11.56-10.08 24.03-15.68 36.58-15.68 21.07,0 41.14,12.21 56.5,35.11 16.81,25.08 24.97,56.67 24.97,89.27 0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75l0-31.02c17.63,0 22.03-16.2 22.03-34.74 0-26.42-6.16-55.74-19.73-76.69-9.63-14.86-22.11-23.94-35.84-23.94-14.85,0-26.8,11.2-40.23,31.17-7.14,10.61-14.47,23.54-22.7,38.13l-9.06,16.05c-18.2,32.27-22.81,39.62-31.91,51.75-15.95,21.24-29.57,29.29-47.5,29.29-21.27,0-34.72-9.21-43.05-23.09-6.8-11.31-10.14-26.15-10.14-43.06z"/>
<path id="Logo1" style="fill:url(#Grad_Logo1)" d="m24.49,37.3c14.24-21.95 34.79-37.3 58.36-37.3 13.65,0 27.22,4.04 41.39,15.61 15.5,12.65 32.02,33.48 52.63,67.81l7.39,12.32c17.84,29.72 27.99,45.01 33.93,52.22 7.64,9.26 12.99,12.02 19.94,12.02 17.63,0 22.03-16.2 22.03-34.74l27.4-.86c0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75-12.8,0-24.14-2.78-36.68-14.61-9.64-9.08-20.91-25.21-29.58-39.71l-25.79-43.08c-12.94-21.62-24.81-37.74-31.68-45.04-7.39-7.85-16.89-17.33-32.05-17.33-12.27,0-22.69,8.61-31.41,21.78z"/>
<path id="Logo2" style="fill:url(#Grad_Logo2)" d="m82.35,31.23c-12.27,0-22.69,8.61-31.41,21.78-12.33,18.61-19.88,46.33-19.88,72.95 0,10.98 2.41,19.41 5.56,24.51l-26.48,17.44c-6.8-11.31-10.14-26.15-10.14-43.06 0-30.75 8.44-62.8 24.49-87.55 14.24-21.95 34.79-37.3 58.36-37.3z"/>
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns:cc="http://creativecommons.org/ns#" width="287.56" height="191">
<desc>Logo of Meta Platforms -- Graphic created by Detmar Owen</desc>
<defs>
<linearGradient id="Grad_Logo1" x1="61" y1="117" x2="259" y2="127" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0064e1" offset="0"/>
<stop style="stop-color:#0064e1" offset="0.4"/>
<stop style="stop-color:#0073ee" offset="0.83"/>
<stop style="stop-color:#0082fb" offset="1"/>
</linearGradient>
<linearGradient id="Grad_Logo2" x1="45" y1="139" x2="45" y2="66" gradientUnits="userSpaceOnUse">
<stop style="stop-color:#0082fb" offset="0"/>
<stop style="stop-color:#0064e0" offset="1"/>
</linearGradient>
</defs>
<path id="Logo0" style="fill:#0081fb" d="m31.06,125.96c0,10.98 2.41,19.41 5.56,24.51 4.13,6.68 10.29,9.51 16.57,9.51 8.1,0 15.51-2.01 29.79-21.76 11.44-15.83 24.92-38.05 33.99-51.98l15.36-23.6c10.67-16.39 23.02-34.61 37.18-46.96 11.56-10.08 24.03-15.68 36.58-15.68 21.07,0 41.14,12.21 56.5,35.11 16.81,25.08 24.97,56.67 24.97,89.27 0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75l0-31.02c17.63,0 22.03-16.2 22.03-34.74 0-26.42-6.16-55.74-19.73-76.69-9.63-14.86-22.11-23.94-35.84-23.94-14.85,0-26.8,11.2-40.23,31.17-7.14,10.61-14.47,23.54-22.7,38.13l-9.06,16.05c-18.2,32.27-22.81,39.62-31.91,51.75-15.95,21.24-29.57,29.29-47.5,29.29-21.27,0-34.72-9.21-43.05-23.09-6.8-11.31-10.14-26.15-10.14-43.06z"/>
<path id="Logo1" style="fill:url(#Grad_Logo1)" d="m24.49,37.3c14.24-21.95 34.79-37.3 58.36-37.3 13.65,0 27.22,4.04 41.39,15.61 15.5,12.65 32.02,33.48 52.63,67.81l7.39,12.32c17.84,29.72 27.99,45.01 33.93,52.22 7.64,9.26 12.99,12.02 19.94,12.02 17.63,0 22.03-16.2 22.03-34.74l27.4-.86c0,19.38-3.82,33.62-10.32,44.87-6.28,10.88-18.52,21.75-39.11,21.75-12.8,0-24.14-2.78-36.68-14.61-9.64-9.08-20.91-25.21-29.58-39.71l-25.79-43.08c-12.94-21.62-24.81-37.74-31.68-45.04-7.39-7.85-16.89-17.33-32.05-17.33-12.27,0-22.69,8.61-31.41,21.78z"/>
<path id="Logo2" style="fill:url(#Grad_Logo2)" d="m82.35,31.23c-12.27,0-22.69,8.61-31.41,21.78-12.33,18.61-19.88,46.33-19.88,72.95 0,10.98 2.41,19.41 5.56,24.51l-26.48,17.44c-6.8-11.31-10.14-26.15-10.14-43.06 0-30.75 8.44-62.8 24.49-87.55 14.24-21.95 34.79-37.3 58.36-37.3z"/>
</svg>
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 2.3 KiB

After

Width:  |  Height:  |  Size: 2.3 KiB

Before After
Before After

12
studio/frontend/src/app/provider.tsx
View file

@ -259,8 +259,16 @@ function TauriWrapper({ children }: { children: ReactNode }) {
<>
{children}
<DownloadManagerPanel />
<WebUpdateBanner enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)} />
<LlamaUpdateBanner enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)} />
<div className="pointer-events-none fixed bottom-4 right-4 z-[9998] flex w-[calc(100vw-2rem)] max-w-[400px] flex-col items-stretch gap-2">
<WebUpdateBanner
positioned={false}
enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)}
/>
<LlamaUpdateBanner
positioned={false}
enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)}
/>
</div>
</>
);
}

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

@ -6,7 +6,7 @@ import { Navbar } from "@/components/navbar";
import { fetchDeviceType, usePlatformStore } from "@/config/env";
import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
import { SettingsDialog, useSettingsDialogStore } from "@/features/settings";
import { useChatRuntimeStore } from "@/features/chat";
import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat";
import { useTrainingUnloadGuard } from "@/features/training";
import { useSidebarPin } from "@/hooks/use-sidebar-pin";
import { useT, type TranslationKey } from "@/i18n";
@ -15,6 +15,7 @@ import {
createRootRoute,
redirect,
useMatches,
useNavigate,
useRouterState,
} from "@tanstack/react-router";
import { AnimatePresence, motion } from "motion/react";
@ -78,6 +79,7 @@ function RootLayout() {
const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname);
const isChatRoute = pathname.startsWith("/chat");
const { pinned, setPinned, togglePinned } = useSidebarPin();
const navigate = useNavigate();
useTrainingUnloadGuard();
@ -107,17 +109,32 @@ function RootLayout() {
if ((e.metaKey || e.ctrlKey) && e.key === ",") {
e.preventDefault();
useSettingsDialogStore.getState().openDialog();
return;
}
// Cmd/Ctrl+Shift+O opens a new chat.
if ((e.metaKey || e.ctrlKey) && e.shiftKey && e.code === "KeyO") {
e.preventDefault();
clearNewChatDraft(); // fresh chat starts empty, no bleed from the last one
const chatRuntime = useChatRuntimeStore.getState();
chatRuntime.setActiveThreadId(null);
chatRuntime.setActiveProjectId(null);
chatRuntime.setIncognito(false);
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
});
}
};
window.addEventListener("keydown", handler);
return () => window.removeEventListener("keydown", handler);
}, []);
}, [navigate]);
useEffect(() => {
if (isChatRoute) return;
const chatRuntime = useChatRuntimeStore.getState();
chatRuntime.setActiveProjectId(null);
chatRuntime.setActiveThreadId(null);
chatRuntime.setIncognito(false);
}, [isChatRoute]);
return (

301
studio/frontend/src/components/app-sidebar.tsx
View file

@ -45,6 +45,7 @@ import { Switch } from "@/components/ui/switch";
import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler";
import { cn } from "@/lib/utils";
import {
Archive03Icon,
ChefHatIcon,
CursorInfo02Icon,
DashboardCircleIcon,
@ -60,11 +61,14 @@ import {
Logout05Icon,
MoreVerticalIcon,
Search01Icon,
PinIcon,
PinOffIcon,
PlusSignIcon,
PowerIcon,
PencilEdit02Icon,
LayoutAlignLeftIcon,
Settings02Icon,
Setting07Icon,
Sun03Icon,
TestTube01Icon,
ZapIcon,
} from "@hugeicons/core-free-icons";
@ -80,10 +84,12 @@ import {
} from "@/components/ui/tooltip";
import { Tooltip as TooltipPrimitive } from "radix-ui";
import { HugeiconsIcon } from "@hugeicons/react";
import { ChevronDown, ChevronsUpDown, MoreHorizontalIcon, Moon, Sun } from "lucide-react";
import { ChevronDown, MoreHorizontalIcon, Moon } from "lucide-react";
import { Link, useNavigate, useRouterState } from "@tanstack/react-router";
import {
archiveChatItem,
ChatSearchDialog,
clearNewChatDraft,
createChatProject,
deleteChatProject,
deleteChatItem,
@ -94,6 +100,8 @@ import {
useChatProjects,
useChatSearchStore,
useChatSidebarItems,
usePinnedChatsStore,
useChatPreferencesStore,
type ProjectRecord,
type SidebarItem,
} from "@/features/chat";
@ -295,11 +303,28 @@ export function AppSidebar() {
enabled: !isStudioRoute,
requireMessages: false,
});
const recentChatItems = useMemo(
() => allChatItems.filter((item) => !item.projectId),
[allChatItems],
const pinnedIds = usePinnedChatsStore((s) => s.pinnedIds);
const togglePinnedChat = usePinnedChatsStore((s) => s.togglePin);
const unpinChat = usePinnedChatsStore((s) => s.unpin);
const confirmDeleteChats = useChatPreferencesStore(
(s) => s.confirmDeleteChats,
);
const chatItems = allChatItems;
const pinnedIdSet = useMemo(() => new Set(pinnedIds), [pinnedIds]);
const recentChatItems = useMemo(
() =>
allChatItems.filter(
(item) => !item.projectId && !pinnedIdSet.has(item.id),
),
[allChatItems, pinnedIdSet],
);
// Pinned chats, in pin order (most recent first).
const pinnedChatItems = useMemo(() => {
const byId = new Map(allChatItems.map((item) => [item.id, item]));
return pinnedIds
.map((id) => byId.get(id))
.filter((item): item is SidebarItem => Boolean(item));
}, [allChatItems, pinnedIds]);
const [pinnedOpen, setPinnedOpen] = useState(true);
const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId);
const activeThreadId = isChatRoute
@ -349,8 +374,12 @@ export function AppSidebar() {
function openNewChat(projectId = activeProjectId) {
if (chatDisabled) return;
clearNewChatDraft();
setActiveThreadId(null);
useChatRuntimeStore.getState().setActiveProjectId(projectId);
// The normal new-chat affordance is always a regular, saved chat --
// only the toolbar toggle starts a temporary one.
useChatRuntimeStore.getState().setIncognito(false);
navigate({ to: "/chat", search: chatSearchForProject(projectId) });
closeMobileIfOpen();
}
@ -374,6 +403,49 @@ export function AppSidebar() {
});
}
// Shared chat delete: same error toast and pin cleanup whether or not the
// confirm dialog is used.
async function deleteChatWithCleanup(item: SidebarItem) {
try {
await handleDeleteThread(item);
unpinChat(item.id);
} catch (err) {
toast.error(translate("shell.toast.failedToDeleteChat"), {
description: err instanceof Error ? err.message : undefined,
});
}
}
async function handleArchiveThread(item: SidebarItem) {
try {
await archiveChatItem(item, activeThreadId, (view) => {
navigate({
to: "/chat",
search: item.projectId
? { project: item.projectId }
: { new: view.newThreadNonce },
});
});
const toastId = toast(
<button
type="button"
onClick={() => {
toast.dismiss(toastId);
useSettingsDialogStore.getState().openArchivedChats();
}}
className="w-full cursor-pointer text-left"
>
You can view archived chats in Settings
</button>,
{ closeButton: true },
);
} catch (err) {
toast.error("Failed to archive chat", {
description: err instanceof Error ? err.message : undefined,
});
}
}
type RenameTarget =
| { kind: "chat"; item: SidebarItem; current: string }
| { kind: "project"; project: ProjectRecord; current: string }
@ -382,6 +454,19 @@ export function AppSidebar() {
null,
);
const [renameDraft, setRenameDraft] = useState("");
// Skips the inline rename input's blur-commit when Enter/Escape already handled it.
const skipRenameBlurRef = useRef(false);
// Optimistic title shown while the debounced sidebar refresh catches up after
// a rename, so the old name does not flash back in.
const [pendingRename, setPendingRename] = useState<{
id: string;
title: string;
} | null>(null);
useEffect(() => {
if (!pendingRename) return;
const match = allChatItems.find((i) => i.id === pendingRename.id);
if (match && match.title === pendingRename.title) setPendingRename(null);
}, [allChatItems, pendingRename]);
const [creatingProject, setCreatingProject] = useState(false);
const [projectNameDraft, setProjectNameDraft] = useState("");
const [projectCreateMoveTarget, setProjectCreateMoveTarget] =
@ -412,9 +497,11 @@ export function AppSidebar() {
if (!target || !renameDirty) return;
setRenamingTarget(null);
if (target.kind === "chat") {
setPendingRename({ id: target.item.id, title: renameTrimmed });
try {
await renameChatItem(target.item, renameTrimmed);
} catch (err) {
setPendingRename(null);
toast.error(translate("shell.toast.failedToRenameChat"), {
description: err instanceof Error ? err.message : undefined,
});
@ -441,6 +528,33 @@ export function AppSidebar() {
}
}
// Inline chat rename commits on Enter or blur, cancels on Escape.
function handleInlineRenameKeyDown(
event: React.KeyboardEvent<HTMLInputElement>,
) {
if (event.key === "Enter") {
event.preventDefault();
skipRenameBlurRef.current = true;
// Commit when changed; otherwise just close, so a no-op Enter does not
// leave the row stuck as an input with its blur suppressed.
if (renameDirty) void commitRename();
else setRenamingTarget(null);
} else if (event.key === "Escape") {
event.preventDefault();
skipRenameBlurRef.current = true;
setRenamingTarget(null);
}
}
function handleInlineRenameBlur() {
if (skipRenameBlurRef.current) {
skipRenameBlurRef.current = false;
return;
}
if (renameDirty) void commitRename();
else setRenamingTarget(null);
}
type DeleteTarget =
| { kind: "chat"; item: SidebarItem }
| { kind: "project"; project: ProjectRecord }
@ -462,13 +576,7 @@ export function AppSidebar() {
target.kind === "project" && deleteProjectFiles;
setConfirmingDelete(null);
if (target.kind === "chat") {
try {
await handleDeleteThread(target.item);
} catch (err) {
toast.error(translate("shell.toast.failedToDeleteChat"), {
description: err instanceof Error ? err.message : undefined,
});
}
await deleteChatWithCleanup(target.item);
return;
}
if (target.kind === "project") {
@ -549,6 +657,7 @@ export function AppSidebar() {
item: SidebarItem,
variant: "project" | "recent",
) {
const isPinned = pinnedIdSet.has(item.id);
const itemClass =
variant === "project"
? "group/project-chat-item relative"
@ -559,13 +668,43 @@ export function AppSidebar() {
: "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto";
const buttonClass = cn(
"sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
// pl-3.5 starts the title at the same x as the Recents label text.
variant === "project" ? "pl-[39px]" : "pl-3.5",
// pl-3 (12px) plus the content's pl-1 (4px) lines the title up with the
// Recents label text at 16px.
variant === "project" ? "pl-[39px]" : "pl-3",
variant === "project"
? "group-hover/project-chat-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8"
: "group-hover/recent-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8",
: isPinned
? // Pinned rows show an extra unpin button on hover, so reserve more room.
"group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8"
: "group-hover/recent-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8",
);
const isRenamingThis =
renamingTarget?.kind === "chat" && renamingTarget.item.id === item.id;
// Inline rename edits the title in place as a rounded pill, no dialog.
if (isRenamingThis) {
return (
<SidebarMenuItem key={item.id} className={itemClass}>
<input
autoFocus
value={renameDraft}
onChange={(event) => setRenameDraft(event.target.value)}
onKeyDown={handleInlineRenameKeyDown}
onBlur={handleInlineRenameBlur}
onFocus={(event) => event.currentTarget.select()}
maxLength={120}
aria-label={translate("shell.dialog.renameChat.placeholder")}
className={cn(
// No pill or box; edit in place as plain highlighted text.
"text-foreground h-[33px] w-full border-0 bg-transparent pr-4 text-[14.5px] leading-[19px] font-medium tracking-nav outline-none",
variant === "project" ? "pl-[39px]" : "pl-3",
)}
/>
</SidebarMenuItem>
);
}
return (
<SidebarMenuItem key={item.id} className={itemClass}>
<SidebarMenuButton
@ -591,7 +730,9 @@ export function AppSidebar() {
closeMobileIfOpen();
}}
>
<span className="truncate">{item.title}</span>
<span className="truncate">
{pendingRename?.id === item.id ? pendingRename.title : item.title}
</span>
</SidebarMenuButton>
<DropdownMenu>
<DropdownMenuTrigger asChild>
@ -610,12 +751,16 @@ export function AppSidebar() {
side="bottom"
align="start"
sideOffset={0}
className="unsloth-plus-menu menu-flat-destructive w-52"
className="unsloth-plus-menu menu-flat-destructive w-56"
>
<DropdownMenuItem onSelect={() => openRenameChat(item)}>
<HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
<span>Rename</span>
</DropdownMenuItem>
<DropdownMenuItem onSelect={() => togglePinnedChat(item.id)}>
<HugeiconsIcon icon={isPinned ? PinOffIcon : PinIcon} strokeWidth={1.75} className="size-icon" />
<span>{isPinned ? "Unpin chat" : "Pin chat"}</span>
</DropdownMenuItem>
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<HugeiconsIcon icon={FolderExportIcon} strokeWidth={1.75} className="size-icon" />
@ -692,15 +837,46 @@ export function AppSidebar() {
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
<DropdownMenuSeparator />
<DropdownMenuItem onSelect={() => void handleArchiveThread(item)}>
<HugeiconsIcon icon={Archive03Icon} strokeWidth={1.75} className="size-icon" />
<span>Archive</span>
</DropdownMenuItem>
<DropdownMenuItem
variant="destructive"
onSelect={() => setConfirmingDelete({ kind: "chat", item })}
onSelect={() =>
confirmDeleteChats
? setConfirmingDelete({ kind: "chat", item })
: void deleteChatWithCleanup(item)
}
>
<HugeiconsIcon icon={Delete02Icon} strokeWidth={1.75} className="size-icon" />
<span>Delete</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
{isPinned ? (
<Tooltip>
<TooltipPrimitive.Trigger asChild>
<button
type="button"
onClick={(e) => {
e.stopPropagation();
unpinChat(item.id);
}}
aria-label="Unpin chat"
className={cn(actionClass, "is-unpin-action")}
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={PinOffIcon} strokeWidth={1.75} className="size-4" />
</span>
</button>
</TooltipPrimitive.Trigger>
<TooltipContent side="bottom" sideOffset={6} className="tooltip-compact">
Unpin
</TooltipContent>
</Tooltip>
) : null}
</SidebarMenuItem>
);
}
@ -937,19 +1113,40 @@ export function AppSidebar() {
</SidebarGroup>
</Collapsible>
{/* Pinned chats: own section above Recents */}
{!isStudioRoute && pinnedChatItems.length > 0 && (
<Collapsible open={pinnedOpen} onOpenChange={setPinnedOpen} asChild>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
<SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
<CollapsibleTrigger className="cursor-pointer flex w-full items-center gap-1 group/sb-collap">
Pinned
<ChevronDown className="size-3.5 opacity-0 transition-[transform,opacity] duration-200 group-hover/sb-collap:opacity-100 group-focus-visible/sb-collap:opacity-100 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg] [[data-state=closed]_&]:opacity-100" />
</CollapsibleTrigger>
</SidebarGroupLabel>
<CollapsibleContent>
<SidebarGroupContent className="pl-1 pr-1.5">
<SidebarMenu>
{pinnedChatItems.map((item) =>
renderChatSidebarItem(item, "recent"),
)}
</SidebarMenu>
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>
</Collapsible>
)}
{!isStudioRoute && (
<Collapsible open={chatOpen} onOpenChange={setChatOpen} asChild>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
<SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
<div className="flex w-full items-center group/sb-collap">
<CollapsibleTrigger className="cursor-pointer flex flex-1 items-center gap-1 min-w-0">
{t("shell.navigation.recents")}
<ChevronDown className="size-3.5 opacity-0 transition-[transform,opacity] duration-200 group-hover/sb-collap:opacity-100 group-focus-visible/sb-collap:opacity-100 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg] [[data-state=closed]_&]:opacity-100" />
</CollapsibleTrigger>
</div>
<CollapsibleTrigger className="cursor-pointer flex w-full items-center gap-1 group/sb-collap">
{t("shell.navigation.recents")}
<ChevronDown className="size-3.5 opacity-0 transition-[transform,opacity] duration-200 group-hover/sb-collap:opacity-100 group-focus-visible/sb-collap:opacity-100 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg] [[data-state=closed]_&]:opacity-100" />
</CollapsibleTrigger>
</SidebarGroupLabel>
<CollapsibleContent>
<SidebarGroupContent className="px-1.5">
<SidebarGroupContent className="pl-1 pr-1.5">
<SidebarMenu>
{recentChatItems.map((item) =>
renderChatSidebarItem(item, "recent"),
@ -971,7 +1168,7 @@ export function AppSidebar() {
</CollapsibleTrigger>
</SidebarGroupLabel>
<CollapsibleContent>
<SidebarGroupContent className="px-2">
<SidebarGroupContent className="pl-1.5 pr-2">
<SidebarMenu>
{runItems.map((run) => {
// Explicit selection wins. Otherwise highlight the active
@ -1092,20 +1289,25 @@ export function AppSidebar() {
<span className="truncate font-heading text-[13.5px] tracking-[0.025em] dark:tracking-[0.04em] font-semibold text-nav-fg">{displayTitle}</span>
<span className="truncate text-[11.5px] tracking-nav text-muted-foreground">Unsloth</span>
</div>
<ChevronsUpDown strokeWidth={1.25} className="ml-auto size-4 text-muted-foreground group-data-[collapsible=icon]:hidden" />
{/* settings cog (replaces the up/down chevron) */}
<HugeiconsIcon
icon={Setting07Icon}
strokeWidth={1.5}
className="ml-auto !size-[18px] text-muted-foreground group-data-[collapsible=icon]:hidden"
/>
</SidebarMenuButton>
</DropdownMenuTrigger>
<DropdownMenuContent
side="top"
align="center"
sideOffset={8}
className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-1.5 py-2.5 font-heading rounded-[20px] border-0"
className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-2.5 py-2.5 font-heading rounded-[20px] border-0"
>
<DropdownMenuGroup>
<DropdownMenuItem
onSelect={() => useSettingsDialogStore.getState().openDialog()}
>
<HugeiconsIcon icon={Settings02Icon} strokeWidth={1.75} className="size-icon" />
<HugeiconsIcon icon={Setting07Icon} strokeWidth={1.75} className="size-icon" />
<span>{t("shell.navigation.settings")}</span>
<DropdownMenuShortcut>⌘,</DropdownMenuShortcut>
</DropdownMenuItem>
@ -1114,7 +1316,7 @@ export function AppSidebar() {
>
<HugeiconsIcon icon={Globe02Icon} strokeWidth={1.75} className="size-[18px]" />
<span>{t("shell.navigation.api")}</span>
<span className="ml-auto rounded-[6px] border border-emerald-500/25 bg-emerald-500/10 px-1.5 py-0.5 text-[10px] leading-none font-semibold text-emerald-700 dark:text-emerald-300">
<span className="ml-auto rounded-full bg-emerald-500/10 px-2 py-1 text-[10px] leading-none font-semibold text-emerald-700 dark:text-emerald-300">
{t("common.new")}
</span>
</DropdownMenuItem>
@ -1122,30 +1324,31 @@ export function AppSidebar() {
ref={anchorRef as React.Ref<HTMLDivElement>}
onSelect={(e) => { e.preventDefault(); toggleTheme(); }}
>
{isDark ? <Sun strokeWidth={1.75} className="size-icon" /> : <Moon strokeWidth={1.75} className="size-icon" />}
{isDark ? <HugeiconsIcon icon={Sun03Icon} strokeWidth={1.75} className="size-icon" /> : <Moon strokeWidth={1.75} className="size-icon" />}
<span>
{isDark
? t("shell.navigation.lightMode")
: t("shell.navigation.darkMode")}
</span>
</DropdownMenuItem>
<DropdownMenuItem
disabled={!getTourId(pathname)}
onSelect={() => {
const tourId = getTourId(pathname);
if (!tourId) return;
window.dispatchEvent(
new CustomEvent(TOUR_OPEN_EVENT, {
detail: { id: tourId },
}),
);
}}
>
<HugeiconsIcon icon={CursorInfo02Icon} strokeWidth={1.75} className="size-icon" />
<span>{t("shell.navigation.guidedTour")}</span>
</DropdownMenuItem>
{getTourId(pathname) && (
<DropdownMenuItem
onSelect={() => {
const tourId = getTourId(pathname);
if (!tourId) return;
window.dispatchEvent(
new CustomEvent(TOUR_OPEN_EVENT, {
detail: { id: tourId },
}),
);
}}
>
<HugeiconsIcon icon={CursorInfo02Icon} strokeWidth={1.75} className="size-icon" />
<span>{t("shell.navigation.guidedTour")}</span>
</DropdownMenuItem>
)}
</DropdownMenuGroup>
<DropdownMenuSeparator className="mx-2.5! my-2.5! h-0! border-t border-border/70 bg-transparent!" />
<DropdownMenuSeparator className="mx-1! my-2.5! h-0! border-t border-border/70 bg-transparent!" />
<DropdownMenuItem
onSelect={() => useSettingsDialogStore.getState().openDialog("about")}
>
@ -1267,7 +1470,7 @@ export function AppSidebar() {
</DialogContent>
</Dialog>
<Dialog
open={renamingTarget !== null}
open={renamingTarget !== null && renamingTarget.kind !== "chat"}
onOpenChange={(open) => {
if (!open) setRenamingTarget(null);
}}

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

@ -8,7 +8,8 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { preprocessLaTeX } from "@/lib/latex";
import { openLink } from "@/lib/open-link";
import { INTERNAL, useAuiState, useMessagePartText } from "@assistant-ui/react";
import { Copy01Icon, Download01Icon, Tick02Icon } from "@hugeicons/core-free-icons";
import { Tick02Icon } from "@/lib/tick-icon";
import { Copy01Icon, Download01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { createMathPlugin } from "@streamdown/math";
import { mermaid } from "@streamdown/mermaid";
@ -284,9 +285,15 @@ function CodeBlockActions({
);
}
// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in
// thread.tsx) and has the artifacts canvas on by default, so a full-HTML answer
// (e.g. a playable game) renders as an interactive card without the global toggle.
function StreamdownBlock(props: BlockProps) {
const shouldCollapseHtmlArtifacts = useChatRuntimeStore(
(state) => state.artifactsEnabled || state.collapseHtmlArtifacts,
(state) =>
state.artifactsEnabled ||
state.collapseHtmlArtifacts ||
state.loadedIsDiffusion,
);
const messageHasRenderableRenderHtmlTool = useAuiState(({ message }) =>
message.parts.some(isRenderableRenderHtmlToolPart),

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

@ -19,6 +19,11 @@ const formatNumber = (n: number): string => {
return n.toLocaleString();
};
const formatRate = (r: number | undefined): string => {
if (r === undefined || !Number.isFinite(r)) return "—";
return `${Math.round(r).toLocaleString()} tok/s`;
};
/**
* Shows streaming stats as a badge with hover tooltip.
* When server timings are available (GGUF), shows prompt eval, generation,
@ -51,6 +56,10 @@ export const MessageTiming: FC<{
st?.cache_n ?? custom?.contextUsage?.cachedTokens ?? 0;
// Anthropic-only cache-write count.
const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0;
// DiffusionGemma reports separately-labelled throughput (no prefill, so no "prompt
// speed"), matching the CLI: in-step parallel, effective (canvas*blocks/wall), and
// output (answer tokens/wall).
const isDiffusion = (st as { diffusion?: boolean } | undefined)?.diffusion === true;
// Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op turns,
// blowing the rate up to Infinity. Require >=1 token, a non-zero decode
@ -93,6 +102,99 @@ export const MessageTiming: FC<{
>
<div className="grid min-w-40 gap-1.5 text-xs">
{st ? (
isDiffusion ? (
<>
{/* DiffusionGemma: honest throughput (no autoregressive prompt speed) */}
{timing.firstTokenTime !== undefined && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">First token</span>
<span className="font-mono tabular-nums">
{formatTimingMs(timing.firstTokenTime)}
</span>
</div>
)}
{st?.diffusion_parallel_tok_s != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Speed (in-step)</span>
<span className="font-mono tabular-nums">
{formatRate(st.diffusion_parallel_tok_s)}
</span>
</div>
)}
{st?.diffusion_effective_tok_s != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Effective</span>
<span className="font-mono tabular-nums">
{formatRate(st.diffusion_effective_tok_s)}
</span>
</div>
)}
{st?.diffusion_output_tok_s != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Output</span>
<span className="font-mono tabular-nums">
{formatRate(st.diffusion_output_tok_s)}
</span>
</div>
)}
{st?.diffusion_steps != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Denoising</span>
<span className="font-mono tabular-nums">
{formatNumber(st.diffusion_steps)} steps
{st?.diffusion_blocks != null
? `, ${formatNumber(st.diffusion_blocks)} block${st.diffusion_blocks === 1 ? "" : "s"}`
: ""}
</span>
</div>
)}
{st?.diffusion_canvas != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Canvas</span>
<span className="font-mono tabular-nums">
{formatNumber(st.diffusion_canvas)} tokens
</span>
</div>
)}
{(st?.diffusion_wall_ms ?? 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">
{formatTimingMs(st.diffusion_wall_ms ?? st.predicted_ms)}
</span>
</div>
)}
{timing.tokenCount !== undefined && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Answer tokens</span>
<span className="font-mono tabular-nums">
{formatNumber(timing.tokenCount)}
</span>
</div>
)}
{(st?.diffusion_prompt_n ?? st?.prompt_n) != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Prompt</span>
<span className="font-mono tabular-nums">
{formatNumber(st.diffusion_prompt_n ?? st.prompt_n)} tokens
</span>
</div>
)}
<div className="my-0.5 border-t border-border/40" />
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Total</span>
<span className="font-mono tabular-nums">
{formatTimingMs(timing.totalStreamTime)}
</span>
</div>
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Chunks</span>
<span className="font-mono tabular-nums">
{timing.totalChunks}
</span>
</div>
</>
) : (
<>
{/* Server-side metrics (GGUF) */}
{st?.prompt_ms != null && (
@ -135,6 +237,30 @@ export const MessageTiming: FC<{
</span>
</div>
)}
{timing.firstTokenTime !== undefined && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">First token</span>
<span className="font-mono tabular-nums">
{formatTimingMs(timing.firstTokenTime)}
</span>
</div>
)}
{st?.diffusion_steps != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Denoising steps</span>
<span className="font-mono tabular-nums">
{formatNumber(st.diffusion_steps)}
</span>
</div>
)}
{st?.diffusion_blocks != null && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Blocks</span>
<span className="font-mono tabular-nums">
{formatNumber(st.diffusion_blocks)}
</span>
</div>
)}
{cacheHits > 0 && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Cache hits</span>
@ -165,6 +291,7 @@ export const MessageTiming: FC<{
</span>
</div>
</>
)
) : (
<>
{/* Client-side metrics (safetensors + external provider fallback) */}

Some files were not shown because too many files have changed in this diff Show more