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

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-06-12 06:35:37 -07:00
commit ecaf3dde2a
234 changed files with 19929 additions and 6549 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

67
install.ps1
View file

@ -1177,14 +1177,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
}
@ -1192,19 +1216,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.
@ -1876,7 +1921,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.2" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -1890,7 +1935,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.2" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -1937,7 +1982,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.2" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
@ -1949,7 +1994,7 @@ shell.Run cmd, 0, False
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.2" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -1977,7 +2022,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.2" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.3" --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)

79
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.2" unsloth-zoo
"unsloth>=2026.6.3" 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.2" unsloth-zoo
"unsloth>=2026.6.3" 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.2" unsloth-zoo
"unsloth>=2026.6.3" 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.2" unsloth-zoo
--upgrade-package unsloth "unsloth>=2026.6.3" 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..."
@ -2672,7 +2689,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.2" --torch-backend=auto
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.3" --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..."

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 -%}

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

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

@ -25,6 +25,7 @@ from utils.hardware import (
get_visible_gpu_count,
)
from core.inference.audio_codecs import AudioCodecManager
from core.inference.runtime_context import runtime_context_length
from io import StringIO
import structlog
from loggers import get_logger
@ -405,6 +406,10 @@ class InferenceBackend:
# Reject CPU/disk offload for audio models too
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
self.models[model_name]["context_length"] = runtime_context_length(
self.models[model_name].get("model"),
max_seq_length,
)
self.active_model_name = model_name
self.loading_models.discard(model_name)
@ -485,6 +490,10 @@ class InferenceBackend:
self.models[model_name]["tokenizer"] = tokenizer
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
self.models[model_name]["context_length"] = runtime_context_length(
self.models[model_name].get("model"),
max_seq_length,
)
self._load_chat_template_info(model_name)

789
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,
)

56
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
@ -227,6 +236,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/mlx_inference.py
View file

@ -7,6 +7,7 @@ instead of torch/transformers for model loading and generation.
import threading
from typing import Optional, Generator
from core.inference.runtime_context import runtime_context_length
from loggers import get_logger
logger = get_logger(__name__)
@ -175,6 +176,7 @@ class MLXInferenceBackend:
"is_audio": False,
"audio_type": None,
"has_audio_input": False,
"context_length": runtime_context_length(self._model, max_seq_length),
}
# Capture chat_template_info so the worker IPC reply ships it back and
# the route layer classifies capabilities like the other paths.

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

@ -727,6 +727,7 @@ class InferenceOrchestrator:
"is_audio": model_info.get("is_audio", False),
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
"context_length": model_info.get("context_length"),
}
# Mirror chat_template_info so routes can classify caps
# without re-entering the subprocess.
@ -860,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,
@ -921,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",

22
studio/backend/core/inference/runtime_context.py Normal file
View file

@ -0,0 +1,22 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Runtime context length helpers shared by inference backends."""
from __future__ import annotations
from typing import Any, Optional
def runtime_context_length(model: Any, fallback: Optional[int] = None) -> Optional[int]:
"""Return the effective context length Unsloth attached to a loaded model."""
for value in (getattr(model, "max_seq_length", None), fallback):
if isinstance(value, bool):
continue
try:
value_int = int(value)
except (TypeError, ValueError):
continue
if value_int > 0:
return value_int
return None

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))

73
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

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

@ -315,6 +315,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
_context_length = _entry.get("context_length")
if _context_length is not None:
model_info["context_length"] = int(_context_length)
except Exception as _ctx_exc:
logger.warning("context_length forward failed: %s", _ctx_exc)
# Forward chat_template_info so the parent can classify capabilities.
try:
_bm = getattr(backend, "models", {}) or {}
@ -881,6 +893,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
name: {
"is_vision": info.get("is_vision", False),
"is_lora": info.get("is_lora", False),
"context_length": info.get("context_length"),
}
for name, info in backend.models.items()
},

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"))

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

@ -40,6 +40,34 @@ logger = get_logger(__name__)
_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.
@ -236,6 +264,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 +338,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 +761,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,

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

@ -1336,6 +1336,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
@ -1515,6 +1541,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")
@ -1693,6 +1739,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
eval_steps = eval_steps_val,
),
)
_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:
@ -1733,7 +1782,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},
@ -1835,26 +1884,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()
@ -2462,7 +2491,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",
@ -2546,6 +2575,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):
@ -3010,20 +3040,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)
@ -3053,17 +3072,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(
{

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

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

58
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)")
@ -171,7 +183,7 @@ class LoadResponse(BaseModel):
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
context_length: Optional[int] = Field(
None, description = "Model's native context length (from GGUF metadata)"
None, description = "Runtime context length in tokens for the loaded model"
)
max_context_length: Optional[int] = Field(
None, description = "Maximum context length currently available on this hardware"
@ -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 = (
@ -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 ─────────────────────
@ -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):

43
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:
@ -338,6 +375,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

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)

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

File diff suppressed because it is too large Load diff

5
studio/backend/routes/llama.py
View file

@ -31,6 +31,7 @@ class LlamaUpdateJob(BaseModel):
from_tag: Optional[str] = None
to_tag: Optional[str] = None
error: Optional[str] = None
progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.")
started_at: Optional[str] = None
finished_at: Optional[str] = None
@ -40,7 +41,9 @@ class LlamaUpdateStatusResponse(BaseModel):
False,
description = "True when the install came from an Unsloth prebuilt (has a marker).",
)
update_available: bool = Field(False, description = "True when installed_tag != latest_tag.")
update_available: bool = Field(
False, description = "True when the latest release is genuinely newer than the install."
)
stale: bool = Field(
False, description = "Update available AND install older than the staleness threshold."
)

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(

16
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", "")
@ -235,6 +246,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 +305,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.

20
studio/backend/run.py
View file

@ -232,6 +232,9 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
public internet. Synchronous so output lands between the banner URLs and the
stop hint. Bounded at ~15s; failures swallowed (verifier failing != Studio
failing). Only meaningful for a wildcard bind."""
global _public_reachable
# Reset to "unknown" each run; set True/False only when the probe decides.
_public_reachable = None
import ipaddress
import json
import time
@ -324,12 +327,14 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
print("", flush = True)
if ok_nodes:
_public_reachable = True
print(
f"{ok_c} Reachability check: {url}/ is reachable from the "
f"public internet ({ok_nodes}/{total} probe nodes connected).{reset}",
flush = True,
)
elif err_nodes:
_public_reachable = False
print(
f"{err_c} Reachability check: {url}/ is NOT reachable from "
f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}",
@ -422,7 +427,9 @@ def _print_cloudflare_line() -> None:
"""Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up.
Reads the module-level URL set by ``run_server``. Prints nothing when the
tunnel is disabled or failed -- failures are silently ignored.
tunnel is disabled or failed -- failures are silently ignored. When the public
reachability probe just failed (``_public_reachable is False``) but the tunnel
is up, reword to point the user at the Cloudflare link as the way in.
"""
if not _cloudflare_url:
return
@ -430,7 +437,10 @@ def _print_cloudflare_line() -> None:
accent = "\033[38;5;150;1m"
reset = "\033[0m"
line = f" Secure link access via Cloudflare: {_cloudflare_url}"
if _public_reachable is False:
line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}"
else:
line = f" Secure link access via Cloudflare: {_cloudflare_url}"
print(f"{accent}{line}{reset}" if stdout_supports_color() else line)
@ -622,6 +632,12 @@ _shutdown_event = None
# None when there is no tunnel (loopback, disabled, or a silently-ignored failure).
_cloudflare_url = None
# Public reachability from the last _verify_global_reachability run, read by the
# Cloudflare banner line. True when the public ip:port probe confirmed reachable,
# False when it confirmed NOT reachable, None when the probe did not run or could
# not decide (timeout, blocked, private address).
_public_reachable = None
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"

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

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

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

@ -13,6 +13,7 @@ import importlib.util
import io
import sys
import tarfile
import types
from pathlib import Path
import pytest
@ -423,3 +424,56 @@ def test_run_server_gates_tunnel_on_wildcard():
source = _RUN_PY.read_text()
assert "_cloudflare_enabled" in source
assert 'host == "0.0.0.0"' in source
def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable):
"""Exec the real _print_cloudflare_line source in isolation (run.py has heavy
deps), with the two module globals injected and startup_banner stubbed."""
src = _RUN_PY.read_text()
tree = ast.parse(src)
func_src = next(
ast.get_source_segment(src, n)
for n in ast.walk(tree)
if isinstance(n, ast.FunctionDef) and n.name == "_print_cloudflare_line"
)
stub = types.ModuleType("startup_banner")
stub.stdout_supports_color = lambda: False
monkeypatch.setitem(sys.modules, "startup_banner", stub)
captured: list[str] = []
ns = {
"_cloudflare_url": cloudflare_url,
"_public_reachable": public_reachable,
"print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)),
}
exec(compile(func_src, "<print_cloudflare_line>", "exec"), ns)
ns["_print_cloudflare_line"]()
return "\n".join(captured)
def test_cloudflare_line_reworded_when_public_unreachable(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = False
)
assert "Use the secure link access via Cloudflare instead: https://x.trycloudflare.com" in out
def test_cloudflare_line_default_wording_when_reachable(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = True
)
assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out
assert "Use the secure link" not in out
def test_cloudflare_line_default_wording_when_unknown(monkeypatch):
# Probe did not run / could not decide -> keep the existing wording.
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = None
)
assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out
assert "Use the secure link" not in out
def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch):
out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False)
assert out == ""

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"

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

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

@ -32,14 +32,19 @@ if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
@pytest.fixture(autouse = True)
def _clear_lemonade_release_cache():
"""Prevent cross-test pollution of the lemonade release lru_cache when
future tests vary the fetch_json mock return value."""
"""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"

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

@ -21,12 +21,25 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
class _NoopLogger:
"""structlog-style logger: every method swallows positional + kwargs.
A stdlib logging.Logger rejects structlog's keyword fields (e.g.
``logger.warning(msg, error=...)``), which leaked into the update module's
error path and failed only when this file's stub loaded first.
"""
def __getattr__(self, _name):
return lambda *a, **k: None
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
_loggers_stub.get_logger = lambda *a, **k: _NoopLogger()
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
_structlog_stub.get_logger = lambda *a, **k: _NoopLogger()
sys.modules.setdefault("structlog", _structlog_stub)
import pytest
@ -51,6 +64,11 @@ def _write_marker(install_dir: Path, **overrides) -> Path:
.replace("+00:00", "Z"),
}
payload.update(overrides)
# The installer always writes `tag` and `release_tag` from the same release
# (a normalized base vs the full release tag), so keep the pair consistent
# when a test overrides only `tag`.
if "tag" in overrides and "release_tag" not in overrides:
payload["release_tag"] = overrides["tag"]
install_dir.mkdir(parents = True, exist_ok = True)
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(payload))
return install_dir / "UNSLOTH_PREBUILT_INFO.json"
@ -303,3 +321,202 @@ def test_format_stale_warning_singular_day():
msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1})
assert "1 day" in msg
assert "1 days" not in msg
# parse_base_build / is_behind.
def test_parse_base_build():
assert fr.parse_base_build("b9596") == 9596
assert fr.parse_base_build(" b9596 ") == 9596
assert fr.parse_base_build("b9596-mix-e6f2453") == 9596 # mix suffix doesn't defeat it
assert fr.parse_base_build("9596") is None
assert fr.parse_base_build("master-abc") is None
assert fr.parse_base_build("") is None
assert fr.parse_base_build(None) is None
@pytest.mark.parametrize(
"installed, latest, expected",
[
(
"b9596-mix-e6f2453",
"b9596-mix-e6f2453",
False,
), # already on the mix latest -> not behind
("b9596", "b9594", False), # latest is an older build -> downgrade guard
("b9596", "b9594-mix-xxx", False), # older mix latest -> still guarded
("b9500", "b9596-mix-e6f2453", True), # newer base -> behind
("b9596-mix-aaa", "b9596-mix-bbb", True), # new mix at same base -> behind
("b9596", "b9596-mix-bbb", True), # clean -> mix at same base -> behind
("b9596-mix-aaa", "b9596", False), # bare base never supersedes a mix install
("b9596", "b9596", False), # identical -> not behind
(" b9596 ", "b9596", False), # whitespace-only diff -> not behind
("master-abc", "master-def", True), # non-bNNNN both -> plain inequality
("master-abc", "master-abc", False),
(None, "b9596", False),
("b9596", None, False),
],
)
def test_is_behind(installed, latest, expected):
assert fr.is_behind(installed, latest) is expected
def test_check_prebuilt_freshness_not_behind_on_mix_latest(monkeypatch, tmp_path):
# Installed the mix latest: marker base tag b9596, full release_tag with sha,
# GitHub latest is that same full tag. Must not report behind (sticky bug).
install_dir = tmp_path / "llama.cpp"
_write_marker(install_dir, tag = "b9596", release_tag = "b9596-mix-e6f2453")
bin_path = _fake_binary(install_dir, layout = "root")
monkeypatch.setattr(
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453"
)
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["behind"] is False
assert info["stale"] is False
def test_check_prebuilt_freshness_downgrade_guard(monkeypatch, tmp_path):
# A lagging latest (older build than installed) must never read as behind/stale.
install_dir = tmp_path / "llama.cpp"
_write_marker(
install_dir,
tag = "b9585",
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 30))
.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: "b9518")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["behind"] is False
assert info["stale"] is False
def test_fetch_latest_release_tag_uses_publish_time(monkeypatch):
# Resolves newest by published_at (like the installer), skips drafts/prereleases,
# and does NOT just take GitHub's first/`/releases/latest` item.
import urllib.request
class _Resp:
def __init__(self, payload):
self._p = json.dumps(payload).encode()
def read(self):
return self._p
def __enter__(self):
return self
def __exit__(self, *a):
return False
payload = [
{
"tag_name": "b9518",
"draft": False,
"prerelease": False,
"published_at": "2026-06-04T21:11:19Z",
},
{
"tag_name": "b9596-mix-e6f2453",
"draft": False,
"prerelease": False,
"published_at": "2026-06-11T22:50:41Z",
},
{
"tag_name": "b9999-draft",
"draft": True,
"prerelease": False,
"published_at": "2026-06-12T00:00:00Z",
},
]
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

171
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:
@ -70,6 +72,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 +1172,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)"})]

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

@ -28,23 +28,75 @@ import utils.llama_cpp_update as upd # noqa: E402
MARKER = "UNSLOTH_PREBUILT_INFO.json"
class _FakeInstallerPopen:
"""Stands in for the streamed installer process in _run_update."""
def __init__(
self,
cmd,
*,
returncode = 0,
lines = None,
on_start = None,
captured_kwargs = None,
**kwargs,
):
if captured_kwargs is not None:
captured_kwargs.update(kwargs)
if on_start is not None:
on_start(list(cmd))
self.returncode = returncode
self.stdout = iter(lines or [])
def wait(self):
return self.returncode
def kill(self):
pass
def _patch_installer_popen(
monkeypatch,
*,
returncode = 0,
lines = None,
on_start = None,
captured_kwargs = None,
):
monkeypatch.setattr(
upd.subprocess,
"Popen",
lambda cmd, **kw: _FakeInstallerPopen(
cmd,
returncode = returncode,
lines = lines,
on_start = on_start,
captured_kwargs = captured_kwargs,
**kw,
),
)
def _write_install(
dir_: Path,
tag: str,
repo: str = "unslothai/llama.cpp",
asset: str | None = None,
release_tag: str | None = None,
) -> str:
"""Create a fake prebuilt install tree and return the llama-server path.
``asset`` is the bundle filename recorded in the marker; omit it to model an
older marker that predates asset-based ROCm forwarding (backward compat)."""
older marker that predates asset-based ROCm forwarding (backward compat).
``release_tag`` is the full release tag (e.g. a ``b9596-mix-<sha>`` mix
build); defaults to ``tag`` for a plain prebuilt."""
bin_dir = dir_ / "build" / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
binary = bin_dir / "llama-server"
binary.write_text("#!/bin/sh\necho stub\n")
marker = {
"tag": tag,
"release_tag": tag,
"release_tag": release_tag or tag,
"published_repo": repo,
"installed_at_utc": "2020-01-01T00:00:00Z",
"bundle_profile": "cuda13-newer",
@ -57,10 +109,13 @@ def _write_install(
@pytest.fixture(autouse = True)
def _clean_state(monkeypatch):
def _clean_state(monkeypatch, tmp_path):
freshness.reset_caches()
upd._reset_job_for_tests()
upd._resolve_memo.clear()
# Isolate the freshness disk cache so the suite never writes the real
# ~/.unsloth cache (the default when storage_roots can't be imported).
monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".freshness_cache")
# Deterministic markerless paths: no host-pinned binary, no custom dir.
monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False)
@ -260,14 +315,15 @@ def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path):
def _fake_run(cmd, **kwargs):
cmd = list(cmd)
# Status polls probe `llama-server --version`; keep the installer argv.
if "--version" in cmd:
return _Proc()
captured["cmd"] = cmd
_write_install(install_dir, "b9585") # installer writes the marker
assert "--version" in cmd # only status polls still use run()
return _Proc()
def _on_start(cmd):
captured["cmd"] = cmd
_write_install(install_dir, "b9585") # installer writes the marker
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
_patch_installer_popen(monkeypatch, on_start = _on_start)
res = upd.start_update()
assert res["started"] is True, res
@ -298,21 +354,27 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
stdout = "installed"
stderr = ""
def _fake_run(cmd, **kwargs):
cmd = list(cmd)
# Status polls probe `llama-server --version`; keep the installer argv.
if "--version" in cmd:
return _Proc()
def _on_start(cmd):
captured["cmd"] = cmd
# Simulate the installer writing a new marker with the latest tag.
_write_install(install_dir, "b9518")
return _Proc()
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
lines = [
"[llama-prebuilt] resolving release\n",
"Downloading llama.zip: 35.0% (12.0 MiB/35.0 MiB) at 9.0 MiB/s\n",
"Downloading llama.zip: 80.0% (28.0 MiB/35.0 MiB) at 9.0 MiB/s\n",
],
on_start = _on_start,
captured_kwargs = popen_kwargs,
)
res = upd.start_update()
assert res["started"] is True
assert res["job"]["from_tag"] == "b9493"
assert res["job"]["progress"] == 0.0
# Wait for the background worker.
deadline = time.time() + 10
@ -328,6 +390,10 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
assert str(install_dir) in captured["cmd"]
assert "--llama-tag" in captured["cmd"] and "latest" in captured["cmd"]
assert "unslothai/llama.cpp" in captured["cmd"]
# Progress lines were parsed and success pins progress at 1.0.
assert job["progress"] == 1.0
# The worker asks the installer for fine-grained progress milestones.
assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path):
@ -335,13 +401,9 @@ def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path):
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
class _Proc:
returncode = 2
stdout = ""
stderr = "boom: network error"
monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc())
_patch_installer_popen(monkeypatch, returncode = 2, lines = ["boom: network error\n"])
res = upd.start_update()
assert res["started"] is True
@ -410,14 +472,15 @@ def _capture_install_cmd(
def _fake_run(cmd, **kwargs):
cmd = list(cmd)
# Status polls probe `llama-server --version`; keep the installer argv.
if "--version" in cmd:
return _Proc()
captured["cmd"] = cmd
_write_install(install_dir, latest, repo = repo, asset = asset)
assert "--version" in cmd # only status polls still use run()
return _Proc()
def _on_start(cmd):
captured["cmd"] = cmd
_write_install(install_dir, latest, repo = repo, asset = asset)
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
_patch_installer_popen(monkeypatch, on_start = _on_start)
res = upd.start_update()
assert res["started"] is True, res
@ -535,18 +598,12 @@ def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path):
seen = {}
class _Proc:
returncode = 0
stdout = "ok"
stderr = ""
def _fake_run(cmd, **kwargs):
def _on_start(cmd):
# The maintenance flag must be set while the installer runs.
seen["flag_during_install"] = backend._llama_update_in_progress
_write_install(install_dir, "b9518")
return _Proc()
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
_patch_installer_popen(monkeypatch, on_start = _on_start)
res = upd.start_update()
assert res["started"] is True
@ -567,16 +624,12 @@ def test_update_clears_maintenance_flag_on_installer_failure(monkeypatch, tmp_pa
binary = _write_install(install_dir, "b9493")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
backend = _FakeBackend()
_inject_backend(monkeypatch, backend)
class _Proc:
returncode = 1
stdout = ""
stderr = "boom"
monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc())
_patch_installer_popen(monkeypatch, returncode = 1, lines = ["boom\n"])
res = upd.start_update()
assert res["started"] is True
@ -606,16 +659,7 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path):
monkeypatch.setitem(sys.modules, "routes", routes_pkg)
monkeypatch.setitem(sys.modules, "routes.inference", inference_mod)
class _Proc:
returncode = 0
stdout = "ok"
stderr = ""
def _fake_run(cmd, **kwargs):
_write_install(install_dir, "b9518")
return _Proc()
monkeypatch.setattr(upd.subprocess, "run", _fake_run)
_patch_installer_popen(monkeypatch, on_start = lambda cmd: _write_install(install_dir, "b9518"))
res = upd.start_update()
assert res["started"] is True
@ -759,3 +803,41 @@ def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path):
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"
# --- mix-tag detection + apply guard (the reported banner bug) ---
def test_status_not_offered_on_mix_latest(monkeypatch, tmp_path):
# Installed the mix latest; GitHub latest is that same full tag -> no banner.
binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453"
)
st = upd.get_update_status()
assert st["update_available"] is False
assert st["installed_tag"] == "b9596"
assert st["latest_tag"] == "b9596-mix-e6f2453"
def test_status_not_offered_when_latest_lags(monkeypatch, tmp_path):
# A lagging latest (older build than installed) must never be offered.
binary = _write_install(tmp_path / "llama.cpp", "b9585")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
st = upd.get_update_status()
assert st["update_available"] is False
def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path):
# A direct POST / stale banner must not reinstall when already on the latest.
binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453"
)
res = upd.start_update()
assert res["started"] is False
assert res["reason"] == "up_to_date"

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

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

@ -19,6 +19,10 @@ from storage import mcp_servers_db
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):

39
studio/backend/tests/test_native_context_length.py
View file

@ -271,6 +271,7 @@ class TestPydanticModels:
def test_load_response_has_field(self):
"""Field exists in LoadResponse.model_fields."""
assert "native_context_length" in LoadResponse.model_fields
assert "context_length" in LoadResponse.model_fields
def test_load_response_defaults_none(self):
"""Omitting native_context_length defaults to None."""
@ -319,6 +320,7 @@ class TestPydanticModels:
def test_status_response_has_field(self):
"""Field exists in InferenceStatusResponse.model_fields."""
assert "native_context_length" in InferenceStatusResponse.model_fields
assert "context_length" in InferenceStatusResponse.model_fields
def test_status_response_has_chat_template_field(self):
"""Status includes chat_template so the UI can rehydrate after refresh."""
@ -347,6 +349,18 @@ class TestPydanticModels:
roundtripped = LoadResponse.model_validate_json(resp.model_dump_json())
assert roundtripped.native_context_length == 131072
def test_context_length_roundtrip(self):
"""Runtime context_length serializes for non-GGUF/hub models."""
resp = LoadResponse(
status = "loaded",
model = "test",
display_name = "Test",
inference = {},
context_length = 8192,
)
roundtripped = LoadResponse.model_validate_json(resp.model_dump_json())
assert roundtripped.context_length == 8192
# =====================================================================
# D. TestRouteCompleteness -- source-level verification
@ -408,6 +422,16 @@ class TestRouteCompleteness:
"native_context_length" not in block
), f"Non-GGUF LoadResponse should not set native_context_length:\n{block[:200]}"
def test_non_gguf_load_responses_set_runtime_context_length(self):
"""Non-GGUF LoadResponse blocks report runtime context_length."""
blocks = self._find_construction_blocks("LoadResponse")
non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b]
assert non_gguf, "Expected at least one non-GGUF LoadResponse block"
for block in non_gguf:
assert (
"context_length" in block
), f"Non-GGUF LoadResponse should set context_length:\n{block[:200]}"
def test_status_path(self):
"""InferenceStatusResponse construction with llama_backend has the field."""
blocks = self._find_construction_blocks("InferenceStatusResponse")
@ -420,6 +444,21 @@ class TestRouteCompleteness:
found
), "No InferenceStatusResponse block with llama_backend has native_context_length"
def test_non_gguf_status_path_reports_runtime_context_length(self):
"""Non-GGUF InferenceStatusResponse reports context_length from model_info."""
blocks = self._find_construction_blocks("InferenceStatusResponse")
found = False
for block in blocks:
if "is_gguf = False" in block and "context_length" in block:
found = True
break
assert found, "No non-GGUF InferenceStatusResponse block with context_length"
def test_openai_models_listing_reports_context_length(self):
"""/v1/models includes context_length when the backend knows it."""
assert 'entry["context_length"]' in self._source
assert 'model_info.get("context_length")' in self._source
# =====================================================================
# E. TestEdgeCases

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

@ -28,11 +28,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 +403,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 +504,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 +520,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 +788,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),
@ -1206,6 +1286,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 +1414,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

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

@ -36,6 +36,7 @@ import json
import httpx
import pytest
from fastapi.responses import JSONResponse
from pydantic import ValidationError
from models.inference import (
@ -46,6 +47,7 @@ from models.inference import (
ResponsesInputMessage,
ResponsesOutputFunctionCall,
ResponsesOutputMessage,
ResponsesOutputReasoning,
ResponsesOutputTextContent,
ResponsesOutputTextPart,
ResponsesRequest,
@ -59,6 +61,7 @@ from routes.inference import (
_chat_tool_calls_to_responses_output,
_normalise_responses_input,
_responses_tool_output_text,
_responses_non_streaming,
_responses_stream,
_translate_responses_tool_choice_to_chat,
_translate_responses_tools_to_chat,
@ -284,6 +287,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
@ -544,6 +600,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 +739,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 +1103,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 +1212,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(
[

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"]

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

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,
}

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

@ -13,6 +13,7 @@ from __future__ import annotations
import json
import os
import re
import time
from datetime import datetime, timezone
from pathlib import Path
@ -104,11 +105,18 @@ def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None:
def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
"""GitHub API call. None on any failure (offline, rate-limited, etc)."""
"""Newest published release tag for `repo`, by publish time.
Resolves "latest" the way install_llama_prebuilt.py does (newest
non-draft/non-prerelease by ``published_at``), NOT via GitHub's
``/releases/latest`` pointer. That pointer sorts by commit date and can lag
behind the build the installer actually installs, so detection and apply
disagreed -- the cause of the downgrade/sticky banner. None on any failure
(offline, rate-limited, etc)."""
import urllib.error
import urllib.request
url = f"https://api.github.com/repos/{repo}/releases/latest"
url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
headers = {
"Accept": "application/vnd.github+json",
"User-Agent": "unsloth-studio-freshness-check",
@ -128,8 +136,21 @@ def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
) as exc:
logger.debug("freshness fetch failed", repo = repo, error = str(exc))
return None
tag = data.get("tag_name")
return tag if isinstance(tag, str) and tag else None
if not isinstance(data, list):
return None
published = [
r
for r in data
if isinstance(r, dict)
and not r.get("draft")
and not r.get("prerelease")
and isinstance(r.get("tag_name"), str)
and r.get("tag_name")
]
if not published:
return None
newest = max(published, key = lambda r: r.get("published_at") or "")
return newest["tag_name"]
def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]:
@ -172,19 +193,58 @@ def _parse_installed_at(value: object) -> Optional[datetime]:
return dt
def parse_base_build(tag: object) -> Optional[int]:
"""Numeric base build from a release tag. Handles both a plain ``bNNNN`` and
a mix-build tag like ``b9596-mix-<sha>`` (anchored at the start, so the mix
suffix doesn't defeat it). None for anything not starting with ``bNNNN``."""
if not isinstance(tag, str):
return None
m = re.match(r"b(\d+)", tag.strip())
return int(m.group(1)) if m else None
def is_behind(installed: Optional[str], latest: Optional[str]) -> bool:
"""Whether `installed` is genuinely behind `latest`, comparing the FULL
release identity (so a mix build can legitimately be the latest) with a
base-build guard so a lagging GitHub /releases/latest can never read as an
update or a downgrade.
- identical tags -> not behind (clears the sticky banner post-update)
- higher base build on `latest` -> behind; lower -> NOT behind (downgrade guard)
- same base build: a different/new mix -> behind, but a bare ``bNNNN`` never
supersedes a mix build (extra PRs) at that base -> not behind
- non-bNNNN tags -> behind (plain inequality, since they already differ)
"""
if not installed or not latest:
return False
installed, latest = installed.strip(), latest.strip()
if installed == latest:
return False
ib, lb = parse_base_build(installed), parse_base_build(latest)
if ib is None or lb is None:
return True
if lb != ib:
return lb > ib
# Same base build, different tags: offer a mix (latest carries a suffix), but
# never offer a bare base over a mix install at the same base.
return latest != f"b{lb}"
def check_prebuilt_freshness(
binary_path: Optional[str],
*,
threshold_days: int = STALENESS_THRESHOLD_DAYS,
now: Optional[datetime] = None,
) -> dict:
"""Returns {has_marker, stale, installed_tag, latest_tag,
"""Returns {has_marker, stale, behind, installed_tag, latest_tag,
installed_at_utc, age_days, published_repo, threshold_days}.
stale = True iff installed != latest AND age >= threshold.
Fails open on missing data (stale stays False)."""
behind = installed genuinely older than latest (see is_behind).
stale = behind AND age >= threshold.
Fails open on missing data (behind/stale stay False)."""
out: dict = {
"has_marker": False,
"stale": False,
"behind": False,
"installed_tag": None,
"latest_tag": None,
"installed_at_utc": None,
@ -196,16 +256,25 @@ def check_prebuilt_freshness(
if not marker:
return out
out["has_marker"] = True
# Display prefers the normalized base ("tag"); comparison below prefers the
# full "release_tag" -- deliberately opposite fallbacks.
out["installed_tag"] = marker.get("tag") or marker.get("release_tag")
out["installed_at_utc"] = marker.get("installed_at_utc")
out["published_repo"] = marker.get("published_repo")
# The marker records both a normalized base tag ("tag", e.g. b9596) and the
# full release tag ("release_tag", e.g. b9596-mix-<sha>). Compare against the
# FULL identity, since GitHub /releases/latest returns the full tag_name --
# comparing the normalized base against the full latest is what produced the
# permanent "downgrade" banner on every mix release.
installed_full = marker.get("release_tag") or marker.get("tag")
repo = out["published_repo"]
if not repo or not out["installed_tag"]:
if not repo or not installed_full:
return out
latest = latest_published_release(repo)
out["latest_tag"] = latest
if not latest or latest == out["installed_tag"]:
out["behind"] = is_behind(installed_full, latest)
if not out["behind"]:
return out
installed_at = _parse_installed_at(out["installed_at_utc"])
@ -232,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)

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

@ -59,10 +59,17 @@ _job: dict = {
"from_tag": None,
"to_tag": None,
"error": None,
"progress": None,
"started_at": None,
"finished_at": None,
}
# Matches the installer's download progress lines, e.g.
# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s".
_PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(")
# The download dominates the update; extract/validate fill the last slice.
_DOWNLOAD_PROGRESS_CEILING = 0.95
def _utcnow() -> str:
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@ -279,9 +286,10 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
freshness = check_prebuilt_freshness(binary)
installed = freshness.get("installed_tag")
latest = freshness.get("latest_tag")
update_available = bool(
freshness.get("has_marker") and installed and latest and installed != latest
)
# `behind` compares the full release identity with a base-build guard, so a
# lagging /releases/latest or a mix-tagged latest can't show a false update
# (see llama_cpp_freshness.is_behind).
update_available = bool(freshness.get("has_marker") and freshness.get("behind"))
with _job_lock:
job = dict(_job)
@ -357,19 +365,58 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
]
cmd.extend(_rocm_install_args(asset))
logger.info("llama update: installing", cmd = " ".join(cmd))
proc = subprocess.run(
# Stream the installer output so download percent lines feed
# job["progress"]; finer milestones via UNSLOTH_PROGRESS_PERCENT_STEP.
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
proc = subprocess.Popen(
cmd,
capture_output = True,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
timeout = _INSTALL_TIMEOUT_SECONDS,
env = env,
)
if proc.returncode != 0:
tail = (proc.stderr or proc.stdout or "").strip()[-1500:]
raise RuntimeError(f"installer exited {proc.returncode}: {tail or 'no output'}")
timed_out = threading.Event()
# New UNSLOTH_PREBUILT_INFO.json is on disk; drop caches so the next
# status read reflects the freshly installed tag.
reset_caches()
def _kill_on_timeout() -> None:
timed_out.set()
proc.kill()
watchdog = threading.Timer(_INSTALL_TIMEOUT_SECONDS, _kill_on_timeout)
watchdog.daemon = True
watchdog.start()
tail_lines: list[str] = []
try:
assert proc.stdout is not None
for line in proc.stdout:
tail_lines.append(line)
if len(tail_lines) > 80:
del tail_lines[0]
m = _PROGRESS_LINE_RE.search(line)
if m is None:
continue
fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING
with _job_lock:
_job["progress"] = max(_job.get("progress") or 0.0, fraction)
returncode = proc.wait()
finally:
watchdog.cancel()
if timed_out.is_set():
raise RuntimeError(f"installer timed out after {_INSTALL_TIMEOUT_SECONDS}s")
if returncode != 0:
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 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
logger.debug("llama update: post-install freshness refresh failed", error = str(exc))
new_marker = read_install_marker(_find_binary())
new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag")
@ -382,6 +429,7 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
),
to_tag = new_tag,
error = None,
progress = 1.0,
finished_at = _utcnow(),
)
logger.info("llama update: success", to_tag = new_tag)
@ -417,7 +465,24 @@ def start_update() -> dict:
"job": get_update_status()["job"],
}
# A job already in flight wins over any freshness re-check below (and skips
# its network call). The final lock block re-checks to close the TOCTOU.
with _job_lock:
if _job["state"] == _JOB_RUNNING:
return {"started": False, "reason": "already_running", "job": dict(_job)}
if marker:
# Mirror the detection guard: a direct POST or a stale banner must not
# start an install when the latest is not actually newer (force a fresh
# check so a stale 24h cache can't wrongly block a real update either).
status = get_update_status(force_refresh = True)
if not status.get("update_available"):
return {
"started": False,
"reason": "up_to_date",
"message": "The installed llama.cpp build is already at the latest prebuilt.",
"job": status["job"],
}
install_dir = _install_dir_for(binary)
repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO
from_tag = marker.get("tag") or marker.get("release_tag")
@ -467,6 +532,7 @@ def start_update() -> dict:
from_tag = from_tag,
to_tag = None,
error = None,
progress = 0.0,
started_at = _utcnow(),
finished_at = None,
)
@ -491,6 +557,7 @@ def _reset_job_for_tests() -> None:
from_tag = None,
to_tag = None,
error = None,
progress = None,
started_at = None,
finished_at = None,
)

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

@ -1371,12 +1371,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 +1386,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

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

@ -43,13 +43,14 @@ from .storage_roots import (
resolve_under_root,
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",
@ -89,6 +90,7 @@ __all__ = [
"resolve_under_root",
"resolve_output_dir",
"resolve_export_dir",
"resolve_export_write_dir",
"resolve_tensorboard_dir",
"resolve_dataset_path",
]

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

@ -6,7 +6,7 @@ from __future__ import annotations
import json
import os
import sys
from pathlib import Path
from pathlib import Path, PurePosixPath, PureWindowsPath
import tempfile
@ -317,6 +317,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 +371,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
@ -373,6 +393,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(),

35
studio/backend/utils/studio_version.py
View file

@ -15,6 +15,7 @@ _DEV_VERSION = "dev"
_GIT_TIMEOUT_SECONDS = 1.0
_STUDIO_TAG_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$")
_GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$")
_GIT_BRANCH_RE = re.compile(r"^[0-9A-Za-z._/-]+$")
_MAX_VERSION_LENGTH = 64
@ -71,6 +72,35 @@ def _exact_git_studio_tag(repo_root: Path) -> str | None:
return tag if is_valid_studio_release_version(tag) else None
def _git_branch(repo_root: Path) -> str | None:
try:
result = subprocess.run(
["git", "rev-parse", "--abbrev-ref", "HEAD"],
cwd = repo_root,
check = False,
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = _GIT_TIMEOUT_SECONDS,
)
except (OSError, subprocess.TimeoutExpired):
return None
if result.returncode != 0:
return None
branch = result.stdout.strip()
# "HEAD" means detached, e.g. a tag or commit checkout.
if (
not branch
or branch == "HEAD"
or len(branch) > _MAX_VERSION_LENGTH
or _GIT_BRANCH_RE.fullmatch(branch) is None
):
return None
return branch
def get_studio_version(repo_root: Path | None = None) -> str:
"""Return the installed Studio release tag for display, or ``dev``.
@ -81,7 +111,10 @@ def get_studio_version(repo_root: Path | None = None) -> str:
if _is_source_checkout(resolved_repo_root):
git_tag = _exact_git_studio_tag(resolved_repo_root)
return git_tag if git_tag is not None else _DEV_VERSION
if git_tag is not None:
return git_tag
branch = _git_branch(resolved_repo_root)
return f"GitHub {branch}" if branch is not None else _DEV_VERSION
stamped_version = _studio_release_build.STUDIO_RELEASE_VERSION
if is_valid_studio_release_version(stamped_version):

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-[340px] 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>
</>
);
}

19
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,11 +109,24 @@ 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);
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
});
}
};
window.addEventListener("keydown", handler);
return () => window.removeEventListener("keydown", handler);
}, []);
}, [navigate]);
useEffect(() => {
if (isChatRoute) return;

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

@ -45,6 +45,8 @@ import { Switch } from "@/components/ui/switch";
import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler";
import { cn } from "@/lib/utils";
import {
Archive01Icon,
ArchiveRestoreIcon,
ChefHatIcon,
CursorInfo02Icon,
DashboardCircleIcon,
@ -60,6 +62,7 @@ import {
Logout05Icon,
MoreVerticalIcon,
Search01Icon,
PlusSignIcon,
PowerIcon,
PencilEdit02Icon,
LayoutAlignLeftIcon,
@ -82,13 +85,16 @@ import { HugeiconsIcon } from "@hugeicons/react";
import { ChevronDown, ChevronsUpDown, MoreHorizontalIcon, Moon, Sun } from "lucide-react";
import { Link, useNavigate, useRouterState } from "@tanstack/react-router";
import {
archiveChatItem,
ChatSearchDialog,
clearNewChatDraft,
createChatProject,
deleteChatProject,
deleteChatItem,
moveChatItemToProject,
renameChatItem,
renameChatProject,
unarchiveChatItem,
useChatRuntimeStore,
useChatProjects,
useChatSearchStore,
@ -290,15 +296,16 @@ export function AppSidebar() {
const activeProjectId = isChatRoute
? ((search.project as string | undefined) ?? null)
: null;
const { items: allChatItems } = useChatSidebarItems({
enabled: !isStudioRoute,
requireMessages: false,
});
const { items: allChatItems, archivedItems: archivedChatItems } =
useChatSidebarItems({
enabled: !isStudioRoute,
requireMessages: false,
});
const recentChatItems = useMemo(
() => allChatItems.filter((item) => !item.projectId),
[allChatItems],
);
const chatItems = allChatItems;
const [archivedOpen, setArchivedOpen] = useState(false);
const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId);
const activeThreadId = isChatRoute
@ -348,6 +355,7 @@ export function AppSidebar() {
function openNewChat(projectId = activeProjectId) {
if (chatDisabled) return;
clearNewChatDraft();
setActiveThreadId(null);
useChatRuntimeStore.getState().setActiveProjectId(projectId);
navigate({ to: "/chat", search: chatSearchForProject(projectId) });
@ -373,6 +381,33 @@ export function AppSidebar() {
});
}
async function handleArchiveThread(item: SidebarItem) {
try {
await archiveChatItem(item, activeThreadId, (view) => {
navigate({
to: "/chat",
search: item.projectId
? { project: item.projectId }
: { new: view.newThreadNonce },
});
});
} catch (err) {
toast.error("Failed to archive chat", {
description: err instanceof Error ? err.message : undefined,
});
}
}
async function handleUnarchiveThread(item: SidebarItem) {
try {
await unarchiveChatItem(item);
} catch (err) {
toast.error("Failed to unarchive chat", {
description: err instanceof Error ? err.message : undefined,
});
}
}
type RenameTarget =
| { kind: "chat"; item: SidebarItem; current: string }
| { kind: "project"; project: ProjectRecord; current: string }
@ -558,7 +593,8 @@ 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",
variant === "project" ? "pl-[39px]" : "pl-3",
// pl-3.5 starts the title at the same x as the Recents label text.
variant === "project" ? "pl-[39px]" : "pl-3.5",
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",
@ -608,7 +644,7 @@ 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" />
@ -690,6 +726,10 @@ export function AppSidebar() {
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
<DropdownMenuItem onSelect={() => void handleArchiveThread(item)}>
<HugeiconsIcon icon={Archive01Icon} strokeWidth={1.75} className="size-icon" />
<span>Archive</span>
</DropdownMenuItem>
<DropdownMenuItem
variant="destructive"
onSelect={() => setConfirmingDelete({ kind: "chat", item })}
@ -837,7 +877,28 @@ export function AppSidebar() {
navigate({ to: "/projects" });
closeMobileIfOpen();
}}
/>
className="group/projects-item relative"
>
<button
type="button"
aria-label="New project"
onClick={(e) => {
e.stopPropagation();
setProjectCreateMoveTarget(null);
setProjectNameDraft("");
setCreatingProject(true);
}}
className="sidebar-row-action group-hover/projects-item:opacity-100 group-hover/projects-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto group-data-[collapsible=icon]:hidden"
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon
icon={PlusSignIcon}
strokeWidth={1.75}
className="size-4"
/>
</span>
</button>
</NavItem>
<NavItem
icon={DashboardCircleIcon}
label={t("shell.navigation.hub")}
@ -926,7 +987,7 @@ export function AppSidebar() {
</div>
</SidebarGroupLabel>
<CollapsibleContent>
<SidebarGroupContent className="px-1.5">
<SidebarGroupContent className="pl-1 pr-1.5">
<SidebarMenu>
{recentChatItems.map((item) =>
renderChatSidebarItem(item, "recent"),
@ -938,6 +999,85 @@ export function AppSidebar() {
</Collapsible>
)}
{/* Archived chats — hidden on Studio + when nothing is archived */}
{!isStudioRoute && archivedChatItems.length > 0 && (
<Collapsible open={archivedOpen} onOpenChange={setArchivedOpen} 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">
Archived
<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">
<SidebarMenu>
{archivedChatItems.map((item) => (
<SidebarMenuItem key={item.id} className="group/archived-item relative">
<SidebarMenuButton
data-testid="archived-thread"
data-thread-type={item.type}
data-thread-id={item.id}
isActive={activeThreadId === item.id}
className="sidebar-nav-btn h-[33px] cursor-pointer rounded-full pl-3.5 pr-4 group-hover/archived-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/archived-item:pr-8 text-[14.5px] leading-[19px] tracking-nav font-medium text-muted-foreground"
onClick={() => {
navigate({
to: "/chat",
search:
item.type === "single"
? { thread: item.id }
: { compare: item.id },
});
closeMobileIfOpen();
}}
>
<span className="truncate">{item.title}</span>
</SidebarMenuButton>
<DropdownMenu>
<DropdownMenuTrigger asChild>
<button
type="button"
onClick={(e) => e.stopPropagation()}
aria-label="Archived chat options"
className="sidebar-row-action group-hover/archived-item:opacity-100 group-hover/archived-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={MoreVerticalIcon} strokeWidth={1.75} className="size-icon" />
</span>
</button>
</DropdownMenuTrigger>
<DropdownMenuContent
side="bottom"
align="start"
sideOffset={0}
className="unsloth-plus-menu menu-flat-destructive w-52"
>
<DropdownMenuItem onSelect={() => openRenameChat(item)}>
<HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
<span>Rename</span>
</DropdownMenuItem>
<DropdownMenuItem onSelect={() => void handleUnarchiveThread(item)}>
<HugeiconsIcon icon={ArchiveRestoreIcon} strokeWidth={1.75} className="size-icon" />
<span>Unarchive</span>
</DropdownMenuItem>
<DropdownMenuItem
variant="destructive"
onSelect={() => setConfirmingDelete({ kind: "chat", item })}
>
<HugeiconsIcon icon={Delete02Icon} strokeWidth={1.75} className="size-icon" />
<span>Delete</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
</SidebarMenuItem>
))}
</SidebarMenu>
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>
</Collapsible>
)}
{isStudioRoute && runItems.length > 0 && !chatOnly && (
<Collapsible open={runsOpen} onOpenChange={setRunsOpen} asChild>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
@ -948,7 +1088,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
@ -1065,7 +1205,7 @@ export function AppSidebar() {
className="!size-[32px]"
/>
</div>
<div className="flex flex-col gap-0.5 leading-tight group-data-[collapsible=icon]:hidden">
<div className="flex flex-col gap-px leading-tight group-data-[collapsible=icon]:hidden">
<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>
@ -1076,7 +1216,7 @@ export function AppSidebar() {
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

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

@ -284,6 +284,9 @@ function CodeBlockActions({
);
}
// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in
// thread.tsx), so it no longer forces HTML into an iframe artifact; it follows the
// same artifact rules as every other model.
function StreamdownBlock(props: BlockProps) {
const shouldCollapseHtmlArtifacts = useChatRuntimeStore(
(state) => state.artifactsEnabled || state.collapseHtmlArtifacts,

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) */}

8
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -189,7 +189,7 @@ function ModelSelectorTrigger({
<HugeiconsIcon
icon={ArrowDown01Icon}
strokeWidth={1.75}
className="relative top-0.5 size-3.5 text-muted-foreground"
className="size-3.5 text-muted-foreground"
/>
</span>
</button>
@ -247,7 +247,7 @@ function ModelSelectorContent({
alignOffset={10}
data-tour={dataTour}
className={cn(
"unsloth-model-selector-menu menu-soft-surface ring-0 w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 p-3",
"unsloth-model-selector-menu menu-soft-surface ring-0 w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 px-3 pt-3 pb-2",
className,
)}
>
@ -307,7 +307,7 @@ function ModelSelectorContent({
)}
{onPickLocalModel ? (
<div className="mt-2 border-t border-border/70 pt-2">
<div className="mt-1.5 border-t border-border/70 pt-1.5">
<button
type="button"
onClick={onPickLocalModel}
@ -320,7 +320,7 @@ function ModelSelectorContent({
</div>
) : null}
{hasSelection && onEject ? (
<div className="mt-2 border-t border-border/70 pt-2">
<div className="mt-1.5 border-t border-border/70 pt-1.5">
<button
type="button"
onClick={onEject}

75
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

@ -39,11 +39,11 @@ import { extractParamLabel } from "@/lib/model-size";
import { cn, formatCompact } from "@/lib/utils";
import type { VramFitStatus } from "@/lib/vram";
import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
import { Add01Icon, Cancel01Icon, Download01Icon, Folder02Icon, Search01Icon } from "@hugeicons/core-free-icons";
import { Add01Icon, Cancel01Icon, Download01Icon, Folder02Icon, Search01Icon, StarIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { FolderBrowser } from "./folder-browser";
import { ModelDeleteAction } from "./model-delete-action";
import { ChevronDownIcon, ChevronRightIcon, StarIcon } from "lucide-react";
import { ChevronDownIcon, ChevronRightIcon } from "lucide-react";
import {
type ReactNode,
useCallback,
@ -63,6 +63,19 @@ function dedupe(values: string[]): string[] {
return [...new Set(values.filter(Boolean))];
}
/** Newest-first by `last_modified` (epoch s), repo_id tie-break. Copies the
* input; treats a missing field as oldest for older-backend compatibility. */
function sortByDownloadRecency<T extends { repo_id: string; last_modified?: number }>(
rows: T[],
): T[] {
return [...rows].sort((a, b) => {
const at = a.last_modified ?? -1;
const bt = b.last_modified ?? -1;
if (at !== bt) return bt - at;
return a.repo_id.localeCompare(b.repo_id);
});
}
/** Lowercase and strip separators for fuzzy search. */
function normalizeForSearch(s: string): string {
return s.toLowerCase().replace(/[\s\-_\.]/g, "");
@ -735,6 +748,33 @@ export function HubModelPicker({
const showHfSection = debouncedQuery.trim().length > 0;
// Newest-first (also covers older backends without `last_modified`).
const sortedCachedGguf = useMemo(
() => sortByDownloadRecency(cachedGguf),
[cachedGguf],
);
const sortedCachedModels = useMemo(
() => sortByDownloadRecency(cachedModels),
[cachedModels],
);
// While searching, filter Downloaded by the query instead of hiding it, so a
// downloaded model the user is searching for stays visible.
const visibleCachedGguf = useMemo(() => {
if (!showHfSection) return sortedCachedGguf;
const q = normalizeForSearch(debouncedQuery.trim());
return sortedCachedGguf.filter((c) => normalizeForSearch(c.repo_id).includes(q));
}, [sortedCachedGguf, showHfSection, debouncedQuery]);
const visibleCachedModels = useMemo(() => {
if (!showHfSection) return sortedCachedModels;
const q = normalizeForSearch(debouncedQuery.trim());
return sortedCachedModels.filter((c) => normalizeForSearch(c.repo_id).includes(q));
}, [sortedCachedModels, showHfSection, debouncedQuery]);
// Non-GGUF cached rows are not shown in chat-only mode, so the empty-state
// logic must use this (not visibleCachedModels) or the picker can go blank.
const visibleCachedModelRows = chatOnly ? [] : visibleCachedModels;
// Recommended models that match the current search query
const filteredRecommendedIds = useMemo(() => {
if (!showHfSection) return [];
@ -765,9 +805,11 @@ export function HubModelPicker({
return results
.map((result) => result.id)
.filter((id) => !recommendedSet.has(id))
// Shown under Downloaded (kept visible while searching); no duplicate.
.filter((id) => !downloadedSet.has(id.toLowerCase()))
.filter((id) => !chatOnly || isKnownGgufRepo(id))
.filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id));
}, [recommendedSet, results, showHfSection, chatOnly, isKnownGgufRepo]);
}, [recommendedSet, downloadedSet, results, showHfSection, chatOnly, isKnownGgufRepo]);
const metricsById = useMemo(
() =>
@ -897,23 +939,28 @@ export function HubModelPicker({
<div ref={scrollRef} className="max-h-64 overflow-y-auto">
<div className="py-1">
{!cachedReady && !showHfSection ? (
{/* First-load spinner only when nothing cached is shown yet. */}
{!cachedReady &&
!showHfSection &&
visibleCachedGguf.length === 0 &&
visibleCachedModelRows.length === 0 ? (
<div className="flex items-center gap-2 px-5 py-3">
<Spinner className="size-3 text-muted-foreground" />
<span className="text-xs text-muted-foreground">
Loading models…
</span>
</div>
) : !showHfSection &&
(cachedGguf.length > 0 ||
(!chatOnly && cachedModels.length > 0)) ? (
) : null}
{/* Downloaded stays visible (filtered) while searching. */}
{visibleCachedGguf.length > 0 || visibleCachedModelRows.length > 0 ? (
<>
<ListLabel
icon={<HugeiconsIcon icon={Download01Icon} className="size-3" />}
collapsed={downloadedCollapsed}
onToggle={() => setDownloadedCollapsed((v) => !v)}
>Downloaded</ListLabel>
{!downloadedCollapsed && cachedGguf.map((c) => (
{!downloadedCollapsed && visibleCachedGguf.map((c) => (
<div key={c.repo_id}>
<ModelRow
label={c.repo_id}
@ -942,8 +989,8 @@ export function HubModelPicker({
)}
</div>
))}
{!downloadedCollapsed && !chatOnly &&
cachedModels.map((c) => (
{!downloadedCollapsed &&
visibleCachedModelRows.map((c) => (
<div key={c.repo_id} className="flex items-center gap-0.5">
<div className="min-w-0 flex-1">
<ModelRow
@ -1231,7 +1278,7 @@ export function HubModelPicker({
{!showHfSection && cachedReady ? (
<>
<ListLabel
icon={<StarIcon className="size-3" />}
icon={<HugeiconsIcon icon={StarIcon} className="size-3" />}
collapsed={recommendedCollapsed}
onToggle={() => setRecommendedCollapsed((v) => !v)}
>Recommended</ListLabel>
@ -1292,7 +1339,7 @@ export function HubModelPicker({
{showHfSection && filteredRecommendedIds.length > 0 ? (
<>
<ListLabel icon={<StarIcon className="size-3" />}>Recommended</ListLabel>
<ListLabel icon={<HugeiconsIcon icon={StarIcon} className="size-3" />}>Recommended</ListLabel>
{filteredRecommendedIds.map((id) => {
const vram = recommendedVramMap.get(id);
return (
@ -1340,7 +1387,9 @@ export function HubModelPicker({
<ListLabel>Hugging Face</ListLabel>
)}
{hfIds.length === 0 && !isLoading ? (
filteredRecommendedIds.length === 0 ? (
filteredRecommendedIds.length === 0 &&
visibleCachedGguf.length === 0 &&
visibleCachedModelRows.length === 0 ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">
No matching models.
</div>

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

@ -19,6 +19,7 @@ import {
thinkEffortAriaLabel,
thinkToggleAriaLabel,
} from "@/components/assistant-ui/think-aria-label";
import { withToolConfirmation } from "@/components/assistant-ui/tool-confirmation-controls";
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { ToolGroup } from "@/components/assistant-ui/tool-group";
import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution";
@ -68,6 +69,14 @@ import { getExternalReasoningCapabilities } from "@/features/chat/provider-capab
import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store";
import {
PLUS_MENU_ORDER,
composerDraftKey,
readComposerDraft,
type PlusMenuItemId,
usePlusMenuPrefsStore,
writeComposerDraft,
} from "@/features/chat";
import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message";
import { ThreadDocumentsBar } from "@/features/rag/components/thread-documents-bar";
import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button";
@ -133,6 +142,7 @@ import {
type KeyboardEvent,
type DragEvent as ReactDragEvent,
type ReactNode,
Fragment,
createContext,
useCallback,
useContext,
@ -936,6 +946,31 @@ const Composer: FC<{
const referenceThreadId = threadId ?? activeThreadId ?? null;
const hasSendableContent =
composerText.trim().length > 0 || hasAttachments || hasPendingAudio;
// Per-thread draft autosave: restore on mount, then mirror composer text
// into localStorage (debounced) so a half-typed message survives a
// navigation or reload. Cleared once empty (i.e. after a send). Setting the
// text even when no draft exists keeps a thread from inheriting the
// previous thread's composer contents.
const draftKey = composerDraftKey(activeThreadId);
const lastDraftKeyRef = useRef(draftKey);
useEffect(() => {
const draft = readComposerDraft(draftKey) ?? "";
const composer = aui.composer();
if (composer.getState().isEditing) {
composer.setText(draft);
}
}, [draftKey, aui]);
useEffect(() => {
// After a thread switch composerText can still hold the previous
// thread's text; skip that cycle so it isn't saved under the new key.
if (lastDraftKeyRef.current !== draftKey) {
lastDraftKeyRef.current = draftKey;
return;
}
const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300);
return () => clearTimeout(t);
}, [composerText, draftKey]);
// Two-row layout shows once the input wraps or a tool is on. Tools can
// pre-select before a model loads, so an active toggle expands it either way.
const composerExpanded =
@ -2092,17 +2127,179 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
const messageCount = useAuiState(({ thread }) => thread.messages.length);
const { startQueue } = useContext(PromptQueueContext);
const plusPins = usePlusMenuPrefsStore((s) => s.pins);
const [recentPrompts, setRecentPrompts] = useState<PromptEntry[]>([]);
const refreshRecentPrompts = useCallback(async () => {
try {
const rows = await listPromptEntries();
setRecentPrompts(
[...rows].sort((a, b) => b.updatedAt - a.updatedAt).slice(0, 3),
);
const byRecent = [...rows].sort((a, b) => b.updatedAt - a.updatedAt);
// Pinned prompts take over the submenu; fall back to the 3 most recent
// when nothing is pinned.
const pinnedIds = usePlusMenuPrefsStore.getState().pinnedPromptIds;
const pinned = byRecent.filter((p) => pinnedIds.includes(p.id));
setRecentPrompts(pinned.length > 0 ? pinned : byRecent.slice(0, 3));
} catch {
}
}, []);
// Adjustable "+" menu items, keyed by id. Pinned ones render at the top
// level; the rest fall into the "More" overflow submenu. The core items
// (photos, web search, code) and "More" itself are always shown and live
// outside this map.
const plusMenuNodes: Record<PlusMenuItemId, ReactNode> = {
chatWithFiles: (
<DropdownMenuItem
disabled={ragDisabled}
className={
ragEnabled && !ragDisabled ? "text-primary font-medium" : undefined
}
onSelect={() => setRagEnabled(!ragEnabled)}
>
<HugeiconsIcon icon={FileDatabaseIcon} strokeWidth={2} />
Chat with Files
{ragEnabled && !ragDisabled ? (
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
) : null}
</DropdownMenuItem>
),
mcp: (
<DropdownMenuItem
disabled={mcpDisabled}
className={
mcpEnabledForChat && !mcpDisabled
? "text-primary font-medium"
: undefined
}
onSelect={() => setMcpEnabledForChat(!mcpEnabledForChat)}
>
<HugeiconsIcon icon={McpServerIcon} strokeWidth={2} />
MCP
{mcpEnabledForChat && !mcpDisabled ? (
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
) : null}
</DropdownMenuItem>
),
savedPrompts: (
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<HugeiconsIcon icon={Bookmark02Icon} strokeWidth={2} />
Saved prompts
</DropdownMenuSubTrigger>
<DropdownMenuSubContent
collisionPadding={16}
className="unsloth-plus-menu w-[208px]"
>
{recentPrompts.map((p) => (
<DropdownMenuItem
key={p.id}
onSelect={() => aui.composer().setText(p.text)}
>
<span className="truncate">{p.name}</span>
</DropdownMenuItem>
))}
{recentPrompts.length > 0 ? <DropdownMenuSeparator /> : null}
<DropdownMenuItem onSelect={() => setPromptStorageOpen(true)}>
All saved prompts…
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
),
compareChat: (
<DropdownMenuItem onSelect={() => startCompare()}>
<Columns2Icon />
Compare chat
</DropdownMenuItem>
),
exportChat: (
<DropdownMenuSub>
<DropdownMenuSubTrigger disabled={!activeThreadId || messageCount === 0}>
<HugeiconsIcon icon={Download01Icon} strokeWidth={2} />
Export chat
</DropdownMenuSubTrigger>
<DropdownMenuSubContent
collisionPadding={16}
className="unsloth-plus-menu w-[208px]"
>
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationRawJsonl(activeThreadId).catch(() =>
toast.error("Export failed."),
);
}}
>
Raw JSONL
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationCsv(activeThreadId).catch(() =>
toast.error("Export failed."),
);
}}
>
CSV
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationShareGPT(activeThreadId).catch(() =>
toast.error("Export failed."),
);
}}
>
ShareGPT JSONL
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
),
canvas: (
<DropdownMenuItem
className={artifactsEnabled ? "text-primary font-medium" : undefined}
onSelect={() => setArtifactsEnabled(!artifactsEnabled)}
>
<HugeiconsIcon icon={PencilRulerIcon} strokeWidth={2} />
Canvas
{artifactsEnabled ? (
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
) : null}
</DropdownMenuItem>
),
projects: (
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<HugeiconsIcon icon={Folder01Icon} strokeWidth={2} />
Projects
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="unsloth-plus-menu w-[232px]">
<DropdownMenuItem onSelect={() => setNewProjectOpen(true)}>
<HugeiconsIcon icon={FolderAddIcon} strokeWidth={2} />
New project
</DropdownMenuItem>
<DropdownMenuLabel>Recents</DropdownMenuLabel>
{recentProjects.length > 0 ? (
recentProjects.map((project) => (
<DropdownMenuItem
key={project.id}
onSelect={() => openProject(project.id)}
>
<HugeiconsIcon icon={Folder01Icon} strokeWidth={2} />
<span className="truncate">{project.name}</span>
</DropdownMenuItem>
))
) : (
<DropdownMenuItem disabled={true}>
No recent projects
</DropdownMenuItem>
)}
</DropdownMenuSubContent>
</DropdownMenuSub>
),
};
const pinnedPlusItems = PLUS_MENU_ORDER.filter((id) => plusPins[id]);
const overflowPlusItems = PLUS_MENU_ORDER.filter((id) => !plusPins[id]);
return (
<>
<PromptStorageDialog
@ -2135,7 +2332,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
align="start"
sideOffset={0}
avoidCollisions={true}
className="unsloth-plus-menu w-[212px]"
className="unsloth-plus-menu w-[244px]"
// Don't refocus the + on close; restored focus showed a stray ring.
onCloseAutoFocus={(event) => event.preventDefault()}
>
@ -2219,165 +2416,22 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
</DropdownMenuItem>
)}
<DropdownMenuSeparator />
<DropdownMenuItem
disabled={ragDisabled}
className={
ragEnabled && !ragDisabled ? "text-primary font-medium" : undefined
}
onSelect={() => setRagEnabled(!ragEnabled)}
>
<HugeiconsIcon icon={FileDatabaseIcon} strokeWidth={2} />
Chat with Files
{ragEnabled && !ragDisabled ? (
<HugeiconsIcon
icon={Tick02Icon}
strokeWidth={2}
className="ml-auto"
/>
) : null}
</DropdownMenuItem>
<DropdownMenuItem
disabled={mcpDisabled}
className={
mcpEnabledForChat && !mcpDisabled
? "text-primary font-medium"
: undefined
}
onSelect={() => setMcpEnabledForChat(!mcpEnabledForChat)}
>
<HugeiconsIcon icon={McpServerIcon} strokeWidth={2} />
MCP
{mcpEnabledForChat && !mcpDisabled ? (
<HugeiconsIcon
icon={Tick02Icon}
strokeWidth={2}
className="ml-auto"
/>
) : null}
</DropdownMenuItem>
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<MoreHorizontalIcon className="size-4" />
More
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="unsloth-plus-menu w-[200px]">
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<HugeiconsIcon icon={Bookmark02Icon} strokeWidth={2} />
Saved prompts
</DropdownMenuSubTrigger>
<DropdownMenuSubContent
collisionPadding={16}
className="unsloth-plus-menu w-[176px]"
>
{recentPrompts.map((p) => (
<DropdownMenuItem
key={p.id}
onSelect={() => aui.composer().setText(p.text)}
>
<span className="truncate">{p.name}</span>
</DropdownMenuItem>
))}
{recentPrompts.length > 0 ? <DropdownMenuSeparator /> : null}
<DropdownMenuItem onSelect={() => setPromptStorageOpen(true)}>
All saved prompts…
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
<DropdownMenuItem onSelect={() => startCompare()}>
<Columns2Icon />
Compare chat
</DropdownMenuItem>
<DropdownMenuSub>
<DropdownMenuSubTrigger
disabled={!activeThreadId || messageCount === 0}
>
<HugeiconsIcon icon={Download01Icon} strokeWidth={2} />
Export chat
</DropdownMenuSubTrigger>
<DropdownMenuSubContent
collisionPadding={16}
className="unsloth-plus-menu w-[176px]"
>
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationRawJsonl(activeThreadId).catch(() =>
toast.error("Export failed."),
);
}}
>
Raw JSONL
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationCsv(activeThreadId).catch(() =>
toast.error("Export failed."),
);
}}
>
CSV
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationShareGPT(activeThreadId).catch(() =>
toast.error("Export failed."),
);
}}
>
ShareGPT JSONL
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
<DropdownMenuItem
className={
artifactsEnabled ? "text-primary font-medium" : undefined
}
onSelect={() => setArtifactsEnabled(!artifactsEnabled)}
>
<HugeiconsIcon icon={PencilRulerIcon} strokeWidth={2} />
Canvas
{artifactsEnabled ? (
<HugeiconsIcon
icon={Tick02Icon}
strokeWidth={2}
className="ml-auto"
/>
) : null}
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
<DropdownMenuSeparator />
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<HugeiconsIcon icon={Folder01Icon} strokeWidth={2} />
Projects
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="unsloth-plus-menu w-[200px]">
<DropdownMenuItem onSelect={() => setNewProjectOpen(true)}>
<HugeiconsIcon icon={FolderAddIcon} strokeWidth={2} />
New project
</DropdownMenuItem>
<DropdownMenuLabel>Recents</DropdownMenuLabel>
{recentProjects.length > 0 ? (
recentProjects.map((project) => (
<DropdownMenuItem
key={project.id}
onSelect={() => openProject(project.id)}
>
<HugeiconsIcon icon={Folder01Icon} strokeWidth={2} />
<span className="truncate">{project.name}</span>
</DropdownMenuItem>
))
) : (
<DropdownMenuItem disabled={true}>
No recent projects
</DropdownMenuItem>
)}
</DropdownMenuSubContent>
</DropdownMenuSub>
{pinnedPlusItems.map((id) => (
<Fragment key={id}>{plusMenuNodes[id]}</Fragment>
))}
{overflowPlusItems.length > 0 ? (
<DropdownMenuSub>
<DropdownMenuSubTrigger>
<MoreHorizontalIcon className="size-4" />
More
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="unsloth-plus-menu w-[232px]">
{overflowPlusItems.map((id) => (
<Fragment key={id}>{plusMenuNodes[id]}</Fragment>
))}
</DropdownMenuSubContent>
</DropdownMenuSub>
) : null}
</DropdownMenuContent>
</DropdownMenu>
<NewProjectDialog
@ -2520,6 +2574,51 @@ const CancelledIndicator: FC = () => {
);
};
const WebSearchToolUIConfirmable = withToolConfirmation(WebSearchToolUI);
const KnowledgeBaseToolUIConfirmable =
withToolConfirmation(KnowledgeBaseToolUI);
const PythonToolUIConfirmable = withToolConfirmation(PythonToolUI);
const TerminalToolUIConfirmable = withToolConfirmation(TerminalToolUI);
const CodeExecutionToolUIConfirmable =
withToolConfirmation(CodeExecutionToolUI);
const ImageGenerationToolUIConfirmable = withToolConfirmation(
ImageGenerationToolUI,
);
const RenderHtmlToolUIConfirmable = withToolConfirmation(RenderHtmlToolUI);
const ToolFallbackConfirmable = withToolConfirmation(ToolFallback);
// Live in-place denoising canvas for DiffusionGemma: while generating, render the
// latest per-step canvas snapshot in the bubble so the user watches the answer resolve
// out of noise. Transient (store-only, cleared on run end), so the finished message
// keeps only the committed markdown.
const DiffusionCanvas: FC = () => {
const isRunning = useAuiState(
({ message }) => message.status?.type === "running",
);
// A non-null canvas is set only by diffusion_frame events (diffusion models only),
// so it is a sufficient gate; loadedIsDiffusion can lag the first frame on a fresh load.
const canvas = useChatRuntimeStore((s) => s.activeDiffusionCanvas);
if (!isRunning || !canvas) {
return null;
}
const stepLabel =
canvas.total > 0 ? `step ${canvas.step + 1}/${canvas.total}` : "denoising";
return (
<div className="aui-diffusion-canvas my-1.5 overflow-hidden rounded-lg border border-primary/20 bg-primary/[0.03]">
<div className="flex items-center gap-2 border-b border-primary/10 px-3 py-1.5 text-[11px] font-medium text-primary/80">
<span className="inline-block size-1.5 animate-pulse rounded-full bg-primary" />
<span>Denoising</span>
<span className="opacity-60">
block {canvas.block + 1} - {stepLabel}
</span>
</div>
<pre className="max-h-[60vh] overflow-auto whitespace-pre-wrap px-3 py-2 font-mono text-[12.5px] leading-relaxed text-foreground/90">
{canvas.text}
</pre>
</div>
);
};
const AssistantMessage: FC = () => {
return (
<MessagePrimitive.Root
@ -2529,6 +2628,7 @@ const AssistantMessage: FC = () => {
<div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
<GeneratingIndicator />
<CancelledIndicator />
<DiffusionCanvas />
<MessagePrimitive.Parts
components={{
Text: MarkdownText,
@ -2538,15 +2638,15 @@ const AssistantMessage: FC = () => {
ToolGroup: ToolGroup,
tools: {
by_name: {
web_search: WebSearchToolUI,
search_knowledge_base: KnowledgeBaseToolUI,
python: PythonToolUI,
terminal: TerminalToolUI,
code_execution: CodeExecutionToolUI,
image_generation: ImageGenerationToolUI,
render_html: RenderHtmlToolUI,
web_search: WebSearchToolUIConfirmable,
search_knowledge_base: KnowledgeBaseToolUIConfirmable,
python: PythonToolUIConfirmable,
terminal: TerminalToolUIConfirmable,
code_execution: CodeExecutionToolUIConfirmable,
image_generation: ImageGenerationToolUIConfirmable,
render_html: RenderHtmlToolUIConfirmable,
},
Fallback: ToolFallback,
Fallback: ToolFallbackConfirmable,
},
}}
/>
@ -2660,10 +2760,10 @@ const AssistantActionBar: FC = () => {
side="bottom"
align="start"
onCloseAutoFocus={(e) => e.preventDefault()}
className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-md [--radius:1.1rem] bg-popover p-1 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-full bg-popover p-1 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
>
<ActionBarPrimitive.ExportMarkdown asChild={true}>
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-sm px-2 py-1.5 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-full px-3 py-1.5 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
<HugeiconsIcon icon={Download01Icon} strokeWidth={1.75} className="size-icon" />
Export as Markdown
</ActionBarMorePrimitive.Item>

157
studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx Normal file
View file

@ -0,0 +1,157 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"use client";
import { Button } from "@/components/ui/button";
import { resolveToolConfirmation } from "@/features/chat/api/chat-api";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import type {
ToolCallMessagePartComponent,
ToolCallMessagePartStatus,
} from "@assistant-ui/react";
import { useCallback, useEffect, useState } from "react";
/**
* Allow / Always allow / Deny controls for a tool call paused awaiting the
* user's confirmation. Rendered alongside every tool card (built-in and
* MCP) so the gate works for all tools, not just the ones using the
* fallback renderer.
*
* A card is "awaiting" only when the adapter registered a backend-gated
* pending call for it (see `toolConfirmations` in the runtime store), so
* non-gated cards -- toggle off, or external-provider tools that already
* ran -- never show controls.
*/
export function ToolConfirmationControls({
toolCallId,
toolName,
result,
status,
}: {
toolCallId?: string;
toolName: string;
result: unknown;
status?: ToolCallMessagePartStatus;
}) {
const confirmation = useChatRuntimeStore((s) =>
toolCallId &&
Object.prototype.hasOwnProperty.call(s.toolConfirmations, toolCallId)
? s.toolConfirmations[toolCallId]
: undefined,
);
const allowToolAlways = useChatRuntimeStore((s) => s.allowToolAlways);
const clearToolConfirmation = useChatRuntimeStore(
(s) => s.clearToolConfirmation,
);
const autoAllowKey = confirmation?.autoAllowKey ?? "";
const autoAllowed = useChatRuntimeStore(
(s) =>
s.alwaysAllowToolsBySession.get(autoAllowKey)?.has(toolName) ?? false,
);
const [decided, setDecided] = useState(false);
const [pending, setPending] = useState<"allow" | "deny" | null>(null);
const [failed, setFailed] = useState(false);
// Still awaiting our decision: a gated pending entry exists, the tool has
// not produced a result, and the card is in its running state.
const awaiting =
confirmation !== undefined &&
result === undefined &&
status?.type === "running";
const showControls = awaiting && !decided;
const resolve = useCallback(
async (decision: "allow" | "deny") => {
if (!toolCallId || !confirmation) return;
setPending(decision);
setFailed(false);
try {
const ok = await resolveToolConfirmation(
confirmation.sessionId,
confirmation.approvalId,
decision,
);
if (ok) {
// Only hide the controls once the backend confirms it matched the
// pending call -- otherwise the generation would stay blocked with
// no way to retry.
setDecided(true);
clearToolConfirmation(toolCallId);
} else {
setFailed(true);
}
} catch {
setFailed(true);
} finally {
setPending(null);
}
},
[toolCallId, confirmation, clearToolConfirmation],
);
// Tools the user marked "Always allow" (this session) approve themselves.
useEffect(() => {
if (showControls && autoAllowed && pending === null && !failed) {
void resolve("allow");
}
}, [showControls, autoAllowed, pending, failed, resolve]);
if (!showControls) return null;
// Auto-approved tools resolve silently unless the post fails.
if (autoAllowed && !failed) return null;
return (
<div className="flex flex-wrap items-center gap-2 pt-1">
<Button
size="xs"
disabled={pending !== null}
onClick={() => void resolve("allow")}
>
Allow
</Button>
<Button
size="xs"
variant="outline"
disabled={pending !== null}
onClick={() => {
if (autoAllowKey) allowToolAlways(autoAllowKey, toolName);
void resolve("allow");
}}
>
Always allow
</Button>
<Button
size="xs"
variant="destructive"
disabled={pending !== null}
onClick={() => void resolve("deny")}
>
Deny
</Button>
{failed ? (
<span className="text-xs text-destructive">
Could not send your decision. Try again.
</span>
) : null}
</div>
);
}
export function withToolConfirmation(
Component: ToolCallMessagePartComponent,
): ToolCallMessagePartComponent {
const WithToolConfirmation: ToolCallMessagePartComponent = (props) => (
<>
<Component {...props} />
<ToolConfirmationControls
toolCallId={props.toolCallId}
toolName={props.toolName}
result={props.result}
status={props.status}
/>
</>
);
return WithToolConfirmation;
}

3
studio/frontend/src/components/assistant-ui/tool-fallback.tsx
View file

@ -325,6 +325,9 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({
result,
status,
}) => {
// Allow/Deny confirmation controls are rendered uniformly for every tool
// card (built-in and fallback) by the `withToolConfirmation` wrapper in
// thread.tsx, so this renderer stays purely presentational.
const isCancelled =
status?.type === "incomplete" && status.reason === "cancelled";

28
studio/frontend/src/components/assistant-ui/tool-group.tsx
View file

@ -9,6 +9,7 @@ import {
type PropsWithChildren,
} from "react";
import { useAuiState } from "@assistant-ui/react";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { ChevronDownIcon } from "lucide-react";
import { Wrench01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@ -216,6 +217,31 @@ const ToolGroupImpl: FC<
(part) => part.type === "tool-call" && part.toolName === "render_html",
),
);
// A blocking allow/deny prompt must never be hidden inside a collapsed
// group, so force the group open while any of its calls awaits confirmation.
const toolConfirmations = useChatRuntimeStore((s) => s.toolConfirmations);
const hasPendingConfirmation = useAuiState(({ message }) =>
message.parts
.slice(startIndex, endIndex + 1)
.some(
(part) =>
part.type === "tool-call" &&
Object.prototype.hasOwnProperty.call(
toolConfirmations,
part.toolCallId,
),
),
);
const messageRunning = useAuiState(
({ message }) => message.status?.type === "running",
);
// Keep the group open once a confirmation forced it open, so answering an
// allow/deny doesn't snap it shut between sequential tool calls. It reverts
// to the default collapsed state once the turn finishes.
const forcedOpenRef = useRef(false);
if (hasPendingConfirmation) forcedOpenRef.current = true;
const forceOpen =
hasPendingConfirmation || (forcedOpenRef.current && messageRunning);
// Render single tool calls and artifacts directly so cards never hide in a
// collapsed group.
@ -224,7 +250,7 @@ const ToolGroupImpl: FC<
}
return (
<ToolGroupRoot>
<ToolGroupRoot open={forceOpen ? true : undefined}>
<ToolGroupTrigger count={toolCount} />
<ToolGroupContent>{children}</ToolGroupContent>
</ToolGroupRoot>

192
studio/frontend/src/components/llama-update-banner.tsx
View file

@ -3,27 +3,100 @@
import { Button } from "@/components/ui/button";
import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check";
import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref";
import { toast } from "@/lib/toast";
import { cn } from "@/lib/utils";
import { AnimatePresence, motion } from "motion/react";
import { type ReactElement, useEffect, useRef } from "react";
import { type ReactElement, useEffect, useRef, useState } from "react";
const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1];
// Backend progress is coarse (5% steps, ~0.9 max) and the extract tail emits no
// signal. Creep toward this cap so the bar keeps moving rather than freezing.
const RUNNING_CAP = 0.95;
// Smoothed 0..1 bar progress: eases toward real `progress`, trickles toward a
// ceiling when idle, animates to 100% when `done`. Resets to 0 on each start.
function useSmoothedProgress(
active: boolean,
progress: number | null,
done: boolean,
): number {
const [display, setDisplay] = useState(0);
const displayRef = useRef(0);
const progressRef = useRef<number | null>(progress);
const doneRef = useRef(done);
progressRef.current = progress;
doneRef.current = done;
useEffect(() => {
if (!active) {
displayRef.current = 0;
setDisplay(0);
return;
}
let raf = 0;
let last = performance.now();
const tick = (now: number) => {
// rAF timestamps can predate the performance.now() captured above, so
// clamp dt at 0 to keep the first frame from stepping backwards.
const dt = Math.max(0, Math.min((now - last) / 1000, 0.1));
last = now;
const current = displayRef.current;
const real = progressRef.current ?? 0;
let target: number;
let speed: number; // approach rate (fraction of remaining gap per second)
if (doneRef.current) {
target = 1;
speed = 5;
} else if (real > current) {
target = real; // catch up to a freshly observed milestone
speed = 4;
} else {
target = RUNNING_CAP; // no signal: creep toward the cap, never frozen
speed = 0.3;
}
const cap = doneRef.current ? 1 : RUNNING_CAP;
const next = Math.min(
current + (target - current) * Math.min(speed * dt, 1),
cap,
);
displayRef.current = next;
setDisplay(next);
if (doneRef.current && next > 0.999) {
return;
}
raf = requestAnimationFrame(tick);
};
raf = requestAnimationFrame(tick);
return () => cancelAnimationFrame(raf);
}, [active]);
return display;
}
interface LlamaUpdateBannerProps {
enabled?: boolean;
// false: fill the parent instead of self-anchoring, so banners can stack in a
// shared container. true (default) keeps standalone desktop mounts working.
positioned?: boolean;
}
/**
* Non-invasive "Update llama.cpp" affordance. Appears bottom-right ~1s after a
* newer prebuilt is detected and stays up until dismissed (click outside / X)
* or updated. Clicking Update swaps the prebuilt in place via POST /api/llama/update.
* newer prebuilt is detected and stays up until the user explicitly acts on it
* (X, Update, or Remind me later). Clicking Update swaps the prebuilt in place
* via POST /api/llama/update. Can be turned off entirely in Settings ->
* General -> Notifications (on by default).
*/
export function LlamaUpdateBanner({
enabled = true,
positioned = true,
}: LlamaUpdateBannerProps): ReactElement | null {
const { status, visible, applying, apply, dismiss } = useLlamaUpdateCheck({
enabled,
});
const showBannerPref = useShowLlamaUpdateBanner();
const { status, visible, applying, apply, dismiss, snooze } =
useLlamaUpdateCheck({
enabled: enabled && showBannerPref,
});
async function handleUpdate() {
const result = await apply();
@ -40,43 +113,36 @@ export function LlamaUpdateBanner({
const show =
visible && status != null && (status.update_available || applying);
const bannerRef = useRef<HTMLDivElement>(null);
// Dismiss when the user clicks anything outside the banner. Kept off while an
// update is applying so the progress stays visible.
useEffect(() => {
if (!show || applying) return;
function onPointerDown(event: PointerEvent) {
if (
bannerRef.current &&
!bannerRef.current.contains(event.target as Node)
) {
dismiss();
}
}
document.addEventListener("pointerdown", onPointerDown, true);
return () =>
document.removeEventListener("pointerdown", onPointerDown, true);
}, [show, applying, dismiss]);
const updateProgress = status?.job.progress ?? null;
const jobSucceeded = status?.job.state === "success";
// Drives the bar so it animates continuously; aria reports the real value.
const displayProgress = useSmoothedProgress(
applying,
updateProgress,
jobSucceeded,
);
return (
<AnimatePresence>
{show ? (
<motion.div
ref={bannerRef}
initial={{ opacity: 0, y: 12, scale: 0.96 }}
animate={{ opacity: 1, y: 0, scale: 1 }}
exit={{ opacity: 0, y: 8, scale: 0.97 }}
transition={{ duration: 0.35, ease: EASE_OUT_QUART }}
className="fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[340px]"
className={cn(
positioned
? "fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[340px]"
: "pointer-events-auto w-full",
)}
data-testid="llama-update-banner"
>
<div className="corner-squircle relative overflow-hidden border border-border/60 bg-background/95 px-4 py-3 shadow-lg backdrop-blur-md">
<div className="relative overflow-hidden rounded-[24px] bg-white px-4 pb-[22px] pl-6 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
{applying ? null : (
<button
type="button"
onClick={dismiss}
className="absolute top-2.5 right-2.5 flex size-5 items-center justify-center rounded-md text-muted-foreground/60 transition-colors hover:bg-muted hover:text-foreground"
className="absolute top-2.5 right-3 flex size-6 items-center justify-center rounded-full text-muted-foreground/60 transition-colors hover:bg-muted hover:text-foreground"
aria-label="Dismiss llama.cpp update notification"
>
<svg
@ -97,32 +163,58 @@ export function LlamaUpdateBanner({
</button>
)}
<div className="flex items-center gap-2 pr-5">
<span className="text-base" aria-hidden="true">
🦥
</span>
<p className="text-sm font-medium text-foreground">
{applying ? "Updating llama.cpp..." : "New llama.cpp prebuilt"}
<div className="min-w-0 pr-6">
<p className="font-heading text-base font-medium text-foreground">
{applying ? "Updating llama.cpp..." : "New llama.cpp version"}
</p>
<p className="mt-0.5 text-xs text-muted-foreground">
{status?.installed_tag ?? "unknown"} &rarr;{" "}
<span className="font-medium text-foreground">
{status?.latest_tag ?? ""}
</span>
</p>
</div>
<p className="mt-0.5 pl-7 text-xs text-muted-foreground">
{status?.installed_tag ?? "unknown"} &rarr;{" "}
<span className="font-medium text-foreground">
{status?.latest_tag ?? ""}
</span>
</p>
<div className="mt-2.5 pl-7">
<Button
size="sm"
className="corner-squircle"
onClick={handleUpdate}
disabled={applying}
data-testid="llama-update-button"
{applying ? (
<div
className="mb-1.5 mt-4 h-1 overflow-hidden rounded-full bg-muted"
role="progressbar"
aria-label="Updating llama.cpp"
aria-valuemin={0}
aria-valuemax={100}
aria-valuenow={
updateProgress != null
? Math.round(updateProgress * 100)
: Math.round(displayProgress * 100)
}
data-testid="llama-update-progress"
>
{applying ? "Updating..." : "Update llama.cpp"}
</Button>
</div>
<div
className="h-full rounded-full bg-primary"
style={{ width: `${Math.max(displayProgress * 100, 2)}%` }}
/>
</div>
) : (
<div className="mt-3 flex items-center gap-2">
<Button
size="sm"
className="h-auto rounded-full px-3.5 py-2 text-[13px]"
onClick={handleUpdate}
data-testid="llama-update-button"
>
Update
</Button>
<Button
size="sm"
variant="ghost"
className="h-auto rounded-full px-2.5 py-2 text-[13px] text-muted-foreground hover:text-foreground"
onClick={snooze}
data-testid="llama-update-snooze-button"
>
Remind me later
</Button>
</div>
)}
</div>
</motion.div>
) : null}

2
studio/frontend/src/components/tauri/startup-screen.tsx
View file

@ -228,7 +228,7 @@ function RepairingContent({
</div>
<div className="mb-10 flex flex-col items-center gap-2">
<TealSpinner />
<p className="text-sm font-bold text-foreground">Updating existing Studio install...</p>
<p className="text-sm font-bold text-foreground">Updating existing Unsloth install...</p>
{latest && (
<p className="max-w-xs text-center text-xs text-muted-foreground">{latest}</p>
)}

190
studio/frontend/src/components/ui/accordion.tsx
View file

@ -1,98 +1,98 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"use client";
import { Accordion as AccordionPrimitive } from "radix-ui";
import type * as React from "react";
import { cn } from "@/lib/utils";
import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
function Accordion({
className,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Root>) {
return (
<AccordionPrimitive.Root
data-slot="accordion"
className={cn(
"overflow-hidden rounded-2xl border flex w-full flex-col",
className,
)}
{...props}
/>
);
}
function AccordionItem({
className,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Item>) {
return (
<AccordionPrimitive.Item
data-slot="accordion-item"
className={cn("data-open:bg-muted/50 not-last:border-b", className)}
{...props}
/>
);
}
function AccordionTrigger({
className,
children,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Trigger>) {
return (
<AccordionPrimitive.Header className="flex">
<AccordionPrimitive.Trigger
data-slot="accordion-trigger"
className={cn(
"**:data-[slot=accordion-trigger-icon]:text-muted-foreground gap-6 p-4 text-left text-sm font-medium hover:underline **:data-[slot=accordion-trigger-icon]:ml-auto **:data-[slot=accordion-trigger-icon]:size-4 group/accordion-trigger relative flex flex-1 items-start justify-between border border-transparent transition-all outline-none disabled:pointer-events-none disabled:opacity-50",
className,
)}
{...props}
>
{children}
<HugeiconsIcon
icon={ArrowDown01Icon}
strokeWidth={2}
data-slot="accordion-trigger-icon"
className="pointer-events-none shrink-0 group-aria-expanded/accordion-trigger:hidden"
/>
<HugeiconsIcon
icon={ArrowUp01Icon}
strokeWidth={2}
data-slot="accordion-trigger-icon"
className="pointer-events-none hidden shrink-0 group-aria-expanded/accordion-trigger:inline"
/>
</AccordionPrimitive.Trigger>
</AccordionPrimitive.Header>
);
}
function AccordionContent({
className,
children,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Content>) {
return (
<AccordionPrimitive.Content
data-slot="accordion-content"
className="data-open:animate-accordion-down data-closed:animate-accordion-up px-4 text-sm overflow-hidden"
{...props}
>
<div
className={cn(
"pt-0 pb-4 [&_a]:hover:text-foreground h-(--radix-accordion-content-height) [&_a]:underline [&_a]:underline-offset-3 [&_p:not(:last-child)]:mb-4",
className,
)}
>
{children}
</div>
</AccordionPrimitive.Content>
);
}
export { Accordion, AccordionItem, AccordionTrigger, AccordionContent };
"use client";
import { Accordion as AccordionPrimitive } from "radix-ui";
import type * as React from "react";
import { cn } from "@/lib/utils";
import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
function Accordion({
className,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Root>) {
return (
<AccordionPrimitive.Root
data-slot="accordion"
className={cn(
"overflow-hidden rounded-2xl border flex w-full flex-col",
className,
)}
{...props}
/>
);
}
function AccordionItem({
className,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Item>) {
return (
<AccordionPrimitive.Item
data-slot="accordion-item"
className={cn("data-open:bg-muted/50 not-last:border-b", className)}
{...props}
/>
);
}
function AccordionTrigger({
className,
children,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Trigger>) {
return (
<AccordionPrimitive.Header className="flex">
<AccordionPrimitive.Trigger
data-slot="accordion-trigger"
className={cn(
"**:data-[slot=accordion-trigger-icon]:text-muted-foreground gap-6 p-4 text-left text-sm font-medium hover:underline **:data-[slot=accordion-trigger-icon]:ml-auto **:data-[slot=accordion-trigger-icon]:size-4 group/accordion-trigger relative flex flex-1 items-start justify-between border border-transparent transition-all outline-none disabled:pointer-events-none disabled:opacity-50",
className,
)}
{...props}
>
{children}
<HugeiconsIcon
icon={ArrowDown01Icon}
strokeWidth={2}
data-slot="accordion-trigger-icon"
className="pointer-events-none shrink-0 group-aria-expanded/accordion-trigger:hidden"
/>
<HugeiconsIcon
icon={ArrowUp01Icon}
strokeWidth={2}
data-slot="accordion-trigger-icon"
className="pointer-events-none hidden shrink-0 group-aria-expanded/accordion-trigger:inline"
/>
</AccordionPrimitive.Trigger>
</AccordionPrimitive.Header>
);
}
function AccordionContent({
className,
children,
...props
}: React.ComponentProps<typeof AccordionPrimitive.Content>) {
return (
<AccordionPrimitive.Content
data-slot="accordion-content"
className="data-open:animate-accordion-down data-closed:animate-accordion-up px-4 text-sm overflow-hidden"
{...props}
>
<div
className={cn(
"pt-0 pb-4 [&_a]:hover:text-foreground h-(--radix-accordion-content-height) [&_a]:underline [&_a]:underline-offset-3 [&_p:not(:last-child)]:mb-4",
className,
)}
>
{children}
</div>
</AccordionPrimitive.Content>
);
}
export { Accordion, AccordionItem, AccordionTrigger, AccordionContent };

368
studio/frontend/src/components/ui/alert-dialog.tsx
View file

@ -1,50 +1,50 @@
// 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 { AlertDialog as AlertDialogPrimitive } from "radix-ui";
import type * as React from "react";
import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";
function AlertDialog({
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Root>) {
return <AlertDialogPrimitive.Root data-slot="alert-dialog" {...props} />;
}
function AlertDialogTrigger({
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Trigger>) {
return (
<AlertDialogPrimitive.Trigger data-slot="alert-dialog-trigger" {...props} />
);
}
function AlertDialogPortal({
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Portal>) {
return (
<AlertDialogPrimitive.Portal data-slot="alert-dialog-portal" {...props} />
);
}
function AlertDialogOverlay({
className,
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Overlay>) {
return (
<AlertDialogPrimitive.Overlay
data-slot="alert-dialog-overlay"
className={cn(
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 bg-black/80 duration-100 supports-backdrop-filter:backdrop-blur-xs fixed inset-0 z-50",
className,
)}
{...props}
/>
);
}
import { AlertDialog as AlertDialogPrimitive } from "radix-ui";
import type * as React from "react";
import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";
function AlertDialog({
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Root>) {
return <AlertDialogPrimitive.Root data-slot="alert-dialog" {...props} />;
}
function AlertDialogTrigger({
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Trigger>) {
return (
<AlertDialogPrimitive.Trigger data-slot="alert-dialog-trigger" {...props} />
);
}
function AlertDialogPortal({
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Portal>) {
return (
<AlertDialogPrimitive.Portal data-slot="alert-dialog-portal" {...props} />
);
}
function AlertDialogOverlay({
className,
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Overlay>) {
return (
<AlertDialogPrimitive.Overlay
data-slot="alert-dialog-overlay"
className={cn(
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 bg-black/80 duration-100 supports-backdrop-filter:backdrop-blur-xs fixed inset-0 z-50",
className,
)}
{...props}
/>
);
}
function AlertDialogContent({
className,
size = "default",
@ -60,143 +60,143 @@ function AlertDialogContent({
<AlertDialogPrimitive.Content
data-slot="alert-dialog-content"
data-size={size}
className={cn(
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 bg-background ring-foreground/5 gap-6 rounded-4xl p-6 ring-1 duration-100 data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-md group/alert-dialog-content fixed top-1/2 left-1/2 z-50 grid w-full -translate-x-1/2 -translate-y-1/2 outline-none",
className,
)}
{...props}
/>
</AlertDialogPortal>
);
}
function AlertDialogHeader({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-dialog-header"
className={cn(
"grid grid-rows-[auto_1fr] place-items-center gap-1.5 text-center has-data-[slot=alert-dialog-media]:grid-rows-[auto_auto_1fr] has-data-[slot=alert-dialog-media]:gap-x-6 sm:group-data-[size=default]/alert-dialog-content:place-items-start sm:group-data-[size=default]/alert-dialog-content:text-left sm:group-data-[size=default]/alert-dialog-content:has-data-[slot=alert-dialog-media]:grid-rows-[auto_1fr]",
className,
)}
{...props}
/>
);
}
function AlertDialogFooter({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-dialog-footer"
className={cn(
"flex flex-col-reverse gap-2 group-data-[size=sm]/alert-dialog-content:grid group-data-[size=sm]/alert-dialog-content:grid-cols-2 sm:flex-row sm:justify-end",
className,
)}
{...props}
/>
);
}
function AlertDialogMedia({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-dialog-media"
className={cn(
"bg-muted mb-2 inline-flex size-16 items-center justify-center rounded-full sm:group-data-[size=default]/alert-dialog-content:row-span-2 *:[svg:not([class*='size-'])]:size-8",
className,
)}
{...props}
/>
);
}
function AlertDialogTitle({
className,
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Title>) {
return (
<AlertDialogPrimitive.Title
data-slot="alert-dialog-title"
className={cn(
"text-lg font-medium sm:group-data-[size=default]/alert-dialog-content:group-has-data-[slot=alert-dialog-media]/alert-dialog-content:col-start-2",
className,
)}
{...props}
/>
);
}
function AlertDialogDescription({
className,
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Description>) {
return (
<AlertDialogPrimitive.Description
data-slot="alert-dialog-description"
className={cn(
"text-muted-foreground *:[a]:hover:text-foreground text-sm text-balance md:text-pretty *:[a]:underline *:[a]:underline-offset-3",
className,
)}
{...props}
/>
);
}
function AlertDialogAction({
className,
variant = "default",
size = "default",
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Action> &
Pick<React.ComponentProps<typeof Button>, "variant" | "size">) {
return (
<Button variant={variant} size={size} asChild>
<AlertDialogPrimitive.Action
data-slot="alert-dialog-action"
className={cn(className)}
{...props}
/>
</Button>
);
}
function AlertDialogCancel({
className,
variant = "outline",
size = "default",
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Cancel> &
Pick<React.ComponentProps<typeof Button>, "variant" | "size">) {
return (
<Button variant={variant} size={size} asChild>
<AlertDialogPrimitive.Cancel
data-slot="alert-dialog-cancel"
className={cn(className)}
{...props}
/>
</Button>
);
}
export {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogMedia,
AlertDialogOverlay,
AlertDialogPortal,
AlertDialogTitle,
AlertDialogTrigger,
};
className={cn(
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 bg-background ring-foreground/5 gap-6 rounded-4xl p-6 ring-1 duration-100 data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-md group/alert-dialog-content fixed top-1/2 left-1/2 z-50 grid w-full -translate-x-1/2 -translate-y-1/2 outline-none",
className,
)}
{...props}
/>
</AlertDialogPortal>
);
}
function AlertDialogHeader({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-dialog-header"
className={cn(
"grid grid-rows-[auto_1fr] place-items-center gap-1.5 text-center has-data-[slot=alert-dialog-media]:grid-rows-[auto_auto_1fr] has-data-[slot=alert-dialog-media]:gap-x-6 sm:group-data-[size=default]/alert-dialog-content:place-items-start sm:group-data-[size=default]/alert-dialog-content:text-left sm:group-data-[size=default]/alert-dialog-content:has-data-[slot=alert-dialog-media]:grid-rows-[auto_1fr]",
className,
)}
{...props}
/>
);
}
function AlertDialogFooter({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-dialog-footer"
className={cn(
"flex flex-col-reverse gap-2 group-data-[size=sm]/alert-dialog-content:grid group-data-[size=sm]/alert-dialog-content:grid-cols-2 sm:flex-row sm:justify-end",
className,
)}
{...props}
/>
);
}
function AlertDialogMedia({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-dialog-media"
className={cn(
"bg-muted mb-2 inline-flex size-16 items-center justify-center rounded-full sm:group-data-[size=default]/alert-dialog-content:row-span-2 *:[svg:not([class*='size-'])]:size-8",
className,
)}
{...props}
/>
);
}
function AlertDialogTitle({
className,
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Title>) {
return (
<AlertDialogPrimitive.Title
data-slot="alert-dialog-title"
className={cn(
"text-lg font-medium sm:group-data-[size=default]/alert-dialog-content:group-has-data-[slot=alert-dialog-media]/alert-dialog-content:col-start-2",
className,
)}
{...props}
/>
);
}
function AlertDialogDescription({
className,
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Description>) {
return (
<AlertDialogPrimitive.Description
data-slot="alert-dialog-description"
className={cn(
"text-muted-foreground *:[a]:hover:text-foreground text-sm text-balance md:text-pretty *:[a]:underline *:[a]:underline-offset-3",
className,
)}
{...props}
/>
);
}
function AlertDialogAction({
className,
variant = "default",
size = "default",
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Action> &
Pick<React.ComponentProps<typeof Button>, "variant" | "size">) {
return (
<Button variant={variant} size={size} asChild>
<AlertDialogPrimitive.Action
data-slot="alert-dialog-action"
className={cn(className)}
{...props}
/>
</Button>
);
}
function AlertDialogCancel({
className,
variant = "outline",
size = "default",
...props
}: React.ComponentProps<typeof AlertDialogPrimitive.Cancel> &
Pick<React.ComponentProps<typeof Button>, "variant" | "size">) {
return (
<Button variant={variant} size={size} asChild>
<AlertDialogPrimitive.Cancel
data-slot="alert-dialog-cancel"
className={cn(className)}
{...props}
/>
</Button>
);
}
export {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogMedia,
AlertDialogOverlay,
AlertDialogPortal,
AlertDialogTitle,
AlertDialogTrigger,
};

152
studio/frontend/src/components/ui/alert.tsx
View file

@ -1,79 +1,79 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { type VariantProps, cva } from "class-variance-authority";
import type * as React from "react";
import { cn } from "@/lib/utils";
const alertVariants = cva(
"grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert",
{
variants: {
variant: {
default: "bg-card text-card-foreground",
destructive:
"text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current",
},
},
defaultVariants: {
variant: "default",
},
},
);
function Alert({
className,
variant,
...props
}: React.ComponentProps<"div"> & VariantProps<typeof alertVariants>) {
return (
<div
data-slot="alert"
role="alert"
className={cn(alertVariants({ variant }), className)}
{...props}
/>
);
}
function AlertTitle({ className, ...props }: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-title"
className={cn(
"font-medium group-has-[>svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3",
className,
)}
{...props}
/>
);
}
function AlertDescription({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-description"
className={cn(
"text-muted-foreground text-sm text-balance md:text-pretty [&_p:not(:last-child)]:mb-4 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3",
className,
)}
{...props}
/>
);
}
function AlertAction({ className, ...props }: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-action"
className={cn("absolute top-2.5 right-3", className)}
{...props}
/>
);
}
export { Alert, AlertTitle, AlertDescription, AlertAction };
import { type VariantProps, cva } from "class-variance-authority";
import type * as React from "react";
import { cn } from "@/lib/utils";
const alertVariants = cva(
"grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert",
{
variants: {
variant: {
default: "bg-card text-card-foreground",
destructive:
"text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current",
},
},
defaultVariants: {
variant: "default",
},
},
);
function Alert({
className,
variant,
...props
}: React.ComponentProps<"div"> & VariantProps<typeof alertVariants>) {
return (
<div
data-slot="alert"
role="alert"
className={cn(alertVariants({ variant }), className)}
{...props}
/>
);
}
function AlertTitle({ className, ...props }: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-title"
className={cn(
"font-medium group-has-[>svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3",
className,
)}
{...props}
/>
);
}
function AlertDescription({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-description"
className={cn(
"text-muted-foreground text-sm text-balance md:text-pretty [&_p:not(:last-child)]:mb-4 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3",
className,
)}
{...props}
/>
);
}
function AlertAction({ className, ...props }: React.ComponentProps<"div">) {
return (
<div
data-slot="alert-action"
className={cn("absolute top-2.5 right-3", className)}
{...props}
/>
);
}
export { Alert, AlertTitle, AlertDescription, AlertAction };

76
studio/frontend/src/components/ui/animated-shiny-text.tsx
View file

@ -1,41 +1,41 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react"
import { cn } from "@/lib/utils"
export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> {
shimmerWidth?: number
}
export const AnimatedShinyText: FC<AnimatedShinyTextProps> = ({
children,
className,
shimmerWidth = 100,
...props
}) => {
return (
<span
style={
{
"--shiny-width": `${shimmerWidth}px`,
} as CSSProperties
}
className={cn(
"mx-auto max-w-md text-neutral-600/70 dark:text-neutral-400/70",
// Shine effect
"animate-shiny-text [background-size:var(--shiny-width)_100%] bg-clip-text [background-position:0_0] bg-no-repeat [transition:background-position_1s_cubic-bezier(.6,.6,0,1)_infinite]",
// Shine gradient
"bg-gradient-to-r from-transparent via-black/80 via-50% to-transparent dark:via-white/80",
className
)}
{...props}
>
{children}
</span>
)
}
import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react"
import { cn } from "@/lib/utils"
export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> {
shimmerWidth?: number
}
export const AnimatedShinyText: FC<AnimatedShinyTextProps> = ({
children,
className,
shimmerWidth = 100,
...props
}) => {
return (
<span
style={
{
"--shiny-width": `${shimmerWidth}px`,
} as CSSProperties
}
className={cn(
"mx-auto max-w-md text-neutral-600/70 dark:text-neutral-400/70",
// Shine effect
"animate-shiny-text [background-size:var(--shiny-width)_100%] bg-clip-text [background-position:0_0] bg-no-repeat [transition:background-position_1s_cubic-bezier(.6,.6,0,1)_infinite]",
// Shine gradient
"bg-gradient-to-r from-transparent via-black/80 via-50% to-transparent dark:via-white/80",
className
)}
{...props}
>
{children}
</span>
)
}

18
studio/frontend/src/components/ui/aspect-ratio.tsx
View file

@ -1,12 +1,12 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { AspectRatio as AspectRatioPrimitive } from "radix-ui";
function AspectRatio({
...props
}: React.ComponentProps<typeof AspectRatioPrimitive.Root>) {
return <AspectRatioPrimitive.Root data-slot="aspect-ratio" {...props} />;
}
export { AspectRatio };
import { AspectRatio as AspectRatioPrimitive } from "radix-ui";
function AspectRatio({
...props
}: React.ComponentProps<typeof AspectRatioPrimitive.Root>) {
return <AspectRatioPrimitive.Root data-slot="aspect-ratio" {...props} />;
}
export { AspectRatio };

220
studio/frontend/src/components/ui/avatar.tsx
View file

@ -1,113 +1,113 @@
// 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 { Avatar as AvatarPrimitive } from "radix-ui";
import type * as React from "react";
import { cn } from "@/lib/utils";
function Avatar({
className,
size = "default",
...props
}: React.ComponentProps<typeof AvatarPrimitive.Root> & {
size?: "default" | "sm" | "lg";
}) {
return (
<AvatarPrimitive.Root
data-slot="avatar"
data-size={size}
className={cn(
"size-8 rounded-full after:rounded-full data-[size=lg]:size-10 data-[size=sm]:size-6 after:border-border group/avatar relative flex shrink-0 select-none after:absolute after:inset-0 after:border after:mix-blend-darken dark:after:mix-blend-lighten",
className,
)}
{...props}
/>
);
}
function AvatarImage({
className,
...props
}: React.ComponentProps<typeof AvatarPrimitive.Image>) {
return (
<AvatarPrimitive.Image
data-slot="avatar-image"
className={cn(
"rounded-full aspect-square size-full object-cover",
className,
)}
{...props}
/>
);
}
function AvatarFallback({
className,
...props
}: React.ComponentProps<typeof AvatarPrimitive.Fallback>) {
return (
<AvatarPrimitive.Fallback
data-slot="avatar-fallback"
className={cn(
"bg-muted text-muted-foreground rounded-full flex size-full items-center justify-center text-sm group-data-[size=sm]/avatar:text-xs",
className,
)}
{...props}
/>
);
}
function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) {
return (
<span
data-slot="avatar-badge"
className={cn(
"bg-primary text-primary-foreground ring-background absolute right-0 bottom-0 z-10 inline-flex items-center justify-center rounded-full bg-blend-color ring-2 select-none",
"group-data-[size=sm]/avatar:size-2 group-data-[size=sm]/avatar:[&>svg]:hidden",
"group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2",
"group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2",
className,
)}
{...props}
/>
);
}
function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) {
return (
<div
data-slot="avatar-group"
className={cn(
"*:data-[slot=avatar]:ring-background group/avatar-group flex -space-x-2 *:data-[slot=avatar]:ring-2",
className,
)}
{...props}
/>
);
}
function AvatarGroupCount({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="avatar-group-count"
className={cn(
"bg-muted text-muted-foreground size-8 rounded-full text-sm group-has-data-[size=lg]/avatar-group:size-10 group-has-data-[size=sm]/avatar-group:size-6 [&>svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2",
className,
)}
{...props}
/>
);
}
export {
Avatar,
AvatarImage,
AvatarFallback,
AvatarGroup,
AvatarGroupCount,
AvatarBadge,
};
import { Avatar as AvatarPrimitive } from "radix-ui";
import type * as React from "react";
import { cn } from "@/lib/utils";
function Avatar({
className,
size = "default",
...props
}: React.ComponentProps<typeof AvatarPrimitive.Root> & {
size?: "default" | "sm" | "lg";
}) {
return (
<AvatarPrimitive.Root
data-slot="avatar"
data-size={size}
className={cn(
"size-8 rounded-full after:rounded-full data-[size=lg]:size-10 data-[size=sm]:size-6 after:border-border group/avatar relative flex shrink-0 select-none after:absolute after:inset-0 after:border after:mix-blend-darken dark:after:mix-blend-lighten",
className,
)}
{...props}
/>
);
}
function AvatarImage({
className,
...props
}: React.ComponentProps<typeof AvatarPrimitive.Image>) {
return (
<AvatarPrimitive.Image
data-slot="avatar-image"
className={cn(
"rounded-full aspect-square size-full object-cover",
className,
)}
{...props}
/>
);
}
function AvatarFallback({
className,
...props
}: React.ComponentProps<typeof AvatarPrimitive.Fallback>) {
return (
<AvatarPrimitive.Fallback
data-slot="avatar-fallback"
className={cn(
"bg-muted text-muted-foreground rounded-full flex size-full items-center justify-center text-sm group-data-[size=sm]/avatar:text-xs",
className,
)}
{...props}
/>
);
}
function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) {
return (
<span
data-slot="avatar-badge"
className={cn(
"bg-primary text-primary-foreground ring-background absolute right-0 bottom-0 z-10 inline-flex items-center justify-center rounded-full bg-blend-color ring-2 select-none",
"group-data-[size=sm]/avatar:size-2 group-data-[size=sm]/avatar:[&>svg]:hidden",
"group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2",
"group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2",
className,
)}
{...props}
/>
);
}
function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) {
return (
<div
data-slot="avatar-group"
className={cn(
"*:data-[slot=avatar]:ring-background group/avatar-group flex -space-x-2 *:data-[slot=avatar]:ring-2",
className,
)}
{...props}
/>
);
}
function AvatarGroupCount({
className,
...props
}: React.ComponentProps<"div">) {
return (
<div
data-slot="avatar-group-count"
className={cn(
"bg-muted text-muted-foreground size-8 rounded-full text-sm group-has-data-[size=lg]/avatar-group:size-10 group-has-data-[size=sm]/avatar-group:size-6 [&>svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2",
className,
)}
{...props}
/>
);
}
export {
Avatar,
AvatarImage,
AvatarFallback,
AvatarGroup,
AvatarGroupCount,
AvatarBadge,
};

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