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

Merge origin/main into image-generation

Browse source 
Resolve the app-sidebar.tsx conflict: main refactored the chat-export dropdown to a
format-based CHAT_EXPORT_OPTIONS + dynamic-import exportConversationByFormat dispatcher,
which the merged body already uses. Keep main's dispatcher and drop the branch's static
export imports; keep TestTubeOutlineIcon imported from the shared @/lib/hugeicons-derived
module (also used by images-page) rather than main's duplicate inline definition. The
branch's Images and Video nav items are preserved.
This commit is contained in:
Daniel Han 2026-07-08 01:57:07 +00:00
commit 27511f30ee
62 changed files with 4429 additions and 286 deletions
Download patch file Download diff file Expand all files Collapse all files

12
.github/scripts/agent-guides-drive.sh vendored
View file

@ -154,7 +154,12 @@ raw_env() { # $1 = var name -> value (one shlex-quote layer stripped)
# writers as a side effect (it writes each agent's relocated session config).
parse_connect() {
local raw="$LOGS_DIR/connect-${AGENT}.txt"
if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
# CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their
# config (which now prompts by default), so the file-edit test opts into auto-approval
# here, the same intent as claude/codex's per-call bypass flags.
local yolo=()
[ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo)
if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
cat_redacted "$raw"
guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero"
fi
@ -394,7 +399,10 @@ case "$MODE" in
T2='Run hello.py with python and show me the exact output.'
# The start.py recipe writers + crosscheck must see the repo; run them
# from the repo root BEFORE cd-ing into the scratch work dir.
# from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw
# gate tool approval through their config (prompting by default), so file-edit
# opts them into auto-approval to run edits/commands headlessly.
case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac
parse_connect
crosscheck_contract
# File-edit needs real tools, so we cannot zero them as in connection.

78
.github/workflows/ossf.yml vendored Normal file
View file

@ -0,0 +1,78 @@
# This workflow uses actions that are not certified by GitHub. They are provided
# by a third-party and are governed by separate terms of service, privacy
# policy, and support documentation.
name: Scorecard supply-chain security
on:
# For Branch-Protection check. Only the default branch is supported. See
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection
branch_protection_rule:
# To guarantee Maintained check is occasionally updated. See
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained
schedule:
- cron: '21 20 * * 0'
push:
branches: [ "main" ]
# Declare default permissions as read only.
permissions: read-all
jobs:
analysis:
name: Scorecard analysis
runs-on: ubuntu-latest
# `publish_results: true` only works when run from the default branch. conditional can be removed if disabled.
if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request'
permissions:
# Needed to upload the results to code-scanning dashboard.
security-events: write
# Needed to publish results and get a badge (see publish_results below).
id-token: write
# Uncomment the permissions below if installing in a private repository.
# contents: read
# actions: read
steps:
- name: "Checkout code"
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- name: "Run analysis"
uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1
with:
results_file: results.sarif
results_format: sarif
# (Optional) "write" PAT token. Uncomment the `repo_token` line below if:
# - you want to enable the Branch-Protection check on a *public* repository, or
# - you are installing Scorecard on a *private* repository
# To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional.
# repo_token: ${{ secrets.SCORECARD_TOKEN }}
# Public repositories:
# - Publish results to OpenSSF REST API for easy access by consumers
# - Allows the repository to include the Scorecard badge.
# - See https://github.com/ossf/scorecard-action#publishing-results.
# For private repositories:
# - `publish_results` will always be set to `false`, regardless
# of the value entered here.
publish_results: true
# (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore
# file_mode: git
# Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
# format to the repository Actions tab.
- name: "Upload artifact"
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: SARIF file
path: results.sarif
retention-days: 5
# Upload the results to GitHub's code scanning dashboard (optional).
# Commenting out will disable upload of results to your repo's Code Scanning dashboard
- name: "Upload to code-scanning"
uses: github/codeql-action/upload-sarif@v3
with:
sarif_file: results.sarif

10
install.ps1
View file

@ -2155,7 +2155,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2169,7 +2169,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2235,7 +2235,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2247,7 +2247,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2275,7 +2275,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --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)

172
install.sh
View file

@ -1483,6 +1483,81 @@ elif [ "$OS" = "macos" ]; then
fi
tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none"
# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in
# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded
# to 10s. Defined here so the reroute below can use it before _run_bounded exists.
_WSL_AMD_GPU_NAME_CACHE=""
_wsl_amd_gpu_name() {
if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then
[ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1
printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0
fi
command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; }
_wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name"
if command -v timeout >/dev/null 2>&1; then
_wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
else
_wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
fi
if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi
_WSL_AMD_GPU_NAME_CACHE="-"; return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04)
# with a 24.04 distro present, re-run the install there and stop; else fall through
# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before
@ -1493,7 +1568,15 @@ _maybe_reroute_strixhalo_to_2404() {
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
[ -e /dev/dxg ] || return 0
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on
# this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
if _has_usable_nvidia_gpu; then return 0; fi
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
# Already ROCm-on-WSL? leave a working GPU alone, whatever the version.
if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then
return 0
@ -1959,61 +2042,6 @@ _has_amd_rocm_gpu() {
return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@ -2327,19 +2355,19 @@ _persist_rocm_wsl_dropin() {
fi
}
# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it.
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
# Leave any already-usable GPU completely alone (NVIDIA, or working ROCm).
if _has_usable_nvidia_gpu; then return 0; fi
# "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the
# generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and
# would skip this bootstrap while the real GPU is still unusable. awk consumes
# all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail.
# Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000,
# the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so
# rocminfo isn't SIGPIPE'd like `grep -q` under pipefail.
_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 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{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
@ -2349,9 +2377,12 @@ _maybe_bootstrap_rocm_wsl() {
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
[ -e /dev/dxg ] || return 0
# Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match
# the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S").
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
command -v bash >/dev/null 2>&1 || return 0
# Fast path: already configured (librocdxg present) but launched from a
@ -2369,7 +2400,8 @@ _maybe_bootstrap_rocm_wsl() {
fi
echo ""
substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN"
_rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU"
substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN"
substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU."
substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)"
@ -2674,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
# 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.
@ -2689,7 +2721,7 @@ if [ "$_MIGRATED" = true ]; then
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-}
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2893,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2911,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
--upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
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..."
@ -2943,7 +2975,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_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --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..."

6
pyproject.toml
View file

@ -73,7 +73,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.1",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -94,7 +94,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.1",
"torchvision",
"unsloth[triton]",
]
@ -579,7 +579,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.1",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

57
scripts/install_rocm_wsl_strixhalo.sh
View file

@ -3,13 +3,14 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# ──────────────────────────────────────────────────────────────────────────────
# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151)
# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX
# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT).
# ──────────────────────────────────────────────────────────────────────────────
# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime
# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG
# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04
# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via
# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies).
# install.sh routes the detected arch to the right ROCm wheels once a runtime exists;
# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg).
# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by
# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the
# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent.
#
# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with
# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once
@ -34,10 +35,12 @@ set -euo pipefail
# ── Tunables (override via env) ──────────────────────────────────────────────
ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install
GFX="gfx1151"
# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200).
# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch.
GFX="${UNSLOTH_WSL_GFX:-}"
LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build
# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test.
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/"
# AMD's wheel index for the (optional) smoke test; resolved after arch detection.
TORCH_INDEX=""
# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs
# torch itself into the real venv right after, so a duplicate download is wasteful.
SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}"
@ -220,12 +223,12 @@ $SUDO ldconfig
say "Persisting ROCm-on-WSL environment"
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
$SUDO tee "$_envfile" >/dev/null <<EOF
# >>> Unsloth ROCm-on-WSL (gfx1151) >>>
# >>> Unsloth ROCm-on-WSL >>>
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${ROCM_DIR}/bin:\${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}"
# <<< Unsloth ROCm-on-WSL (gfx1151) <<<
# <<< Unsloth ROCm-on-WSL <<<
EOF
# also drop into ~/.bashrc for interactive shells
if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then
@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}"
# ── Step 5: verify the runtime enumerates the GPU ────────────────────────────
say "Verifying rocminfo sees ${GFX}"
say "Verifying rocminfo enumerates the GPU over DXG"
# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs
# rocminfo on first match, which under `set -o pipefail` turns a successful match
# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a
# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass.
# into a pipeline failure.
_rocminfo_out="$(rocminfo 2>/dev/null || true)"
if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then
# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU
# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch.
_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)"
if [ -z "$_detected_gfx" ]; then
printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true
die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
fi
# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under
# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt.
if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then
die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'."
fi
GFX="${GFX:-$_detected_gfx}"
# Display-only summary: best-effort (|| true) so head's early pipe-close under
# `set -o pipefail` can't fail the bootstrap after verification already passed.
printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true
note "ROCm-on-WSL runtime is live for ${GFX}."
# ── Step 6 (optional): torch smoke test from the gfx1151 index ───────────────
# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ───────
if [ "$SMOKE_TEST" = "1" ]; then
say "Smoke-testing PyTorch on ${GFX} (throwaway venv)"
# Map the detected arch to AMD's repo.amd.com wheel family index.
case "$GFX" in
gfx1200|gfx1201) _fam="gfx120X-all" ;;
gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;;
*) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index
esac
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/"
_venv="${HOME}/.unsloth/rocm-smoketest"
rm -rf "$_venv"; python3 -m venv "$_venv"
"$_venv/bin/pip" install --quiet --upgrade pip
# gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py
# AMD arch index is primary (torch + triton); PyPI only an extra for pure-py
# deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch.
"$_venv/bin/pip" install --index-url "$TORCH_INDEX" \
--extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \
die "torch install from ${TORCH_INDEX} failed."
# WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib.
_tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)"
[ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true
"$_venv/bin/python" - <<'PY'
import torch
ok = torch.cuda.is_available()

9
studio/backend/assets/configs/inference_defaults.json
View file

@ -235,6 +235,13 @@
"min_p": 0.1,
"repetition_penalty": 1.0
},
"deepseek-v4": {
"temperature": 1.0,
"top_p": 1.0,
"top_k": -1,
"min_p": 0.0,
"repetition_penalty": 1.0
},
"deepseek-r1": {
"temperature": 0.6,
"top_p": 0.95,
@ -394,7 +401,7 @@
"phi-4", "phi-3",
"mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral",
"devstral", "pixtral",
"deepseek-r1", "deepseek-v3", "deepseek-ocr",
"deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr",
"glm-5", "glm-4",
"nemotron",
"minimax-m2.7", "minimax-m2.5", "minimax",

28
studio/backend/core/inference/anthropic_compat.py
View file

@ -258,6 +258,10 @@ class AnthropicStreamEmitter:
self._open_tool_use_id: Optional[str] = None
self._open_tool_args_sent: bool = False
self._prev_text: str = ""
# Net <think> minus </think> in the text emitted to the client. Tracked
# from emitted deltas (not _prev_text, which a final bare shrink clobbers)
# so an unclosed reasoning-only block can be balanced before close.
self._open_think_tags: int = 0
self._usage: dict = {}
def start(
@ -317,6 +321,7 @@ class AnthropicStreamEmitter:
"""Close any open block and emit message_delta + message_stop."""
events = []
if self._text_block_open or self._open_tool_call_id is not None:
events.extend(self._close_open_think())
events.append(self._close_block())
self._open_tool_call_id = None
self._open_tool_use_id = None
@ -344,12 +349,33 @@ class AnthropicStreamEmitter:
)
return events
def _close_open_think(self) -> list[str]:
"""Emit a ``</think>`` delta when the streamed text left a ``<think>``
open. This emitter diffs cumulative snapshots and drops the generator's
final bare shrink, so a reasoning-only reply would otherwise end on an
unclosed tag. Mirrors the chat route's reasoning extractor, which closes
the block on finish; balances the block before it is closed."""
if not self._text_block_open or self._open_think_tags <= 0:
return []
self._open_think_tags = 0
return [
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {"type": "text_delta", "text": "</think>"},
},
)
]
def _handle_content(self, event: dict) -> list[str]:
cumulative = event.get("text", "")
new_text = cumulative[len(self._prev_text) :]
self._prev_text = cumulative
if not new_text:
return []
self._open_think_tags += new_text.count("<think>") - new_text.count("</think>")
if not self._text_block_open:
events = self._open_text_block()
else:
@ -374,6 +400,7 @@ class AnthropicStreamEmitter:
events = []
if self._text_block_open:
events.extend(self._close_open_think())
events.append(self._close_block())
# Defensive: close a stale open tool_use block before starting another.
elif self._open_tool_call_id is not None:
@ -452,6 +479,7 @@ class AnthropicStreamEmitter:
events.extend(self._open_text_block())
# Reset text tracking for the next synthesis turn
self._prev_text = ""
self._open_think_tags = 0
return events
def _open_text_block(self) -> list[str]:

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

@ -8,6 +8,7 @@ import utils.hardware.hardware as hw
DEFAULT_MODELS_GGUF = [
"unsloth/Qwen3.6-27B-MTP-GGUF",
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
"unsloth/DeepSeek-V4-Flash-GGUF",
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",
@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [
DEFAULT_MODELS_STANDARD = [
"unsloth/Qwen3.6-27B-MTP-GGUF",
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
"unsloth/DeepSeek-V4-Flash-GGUF",
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",

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

@ -686,6 +686,16 @@ def detect_reasoning_flags(
else []
)
if effort_levels:
# DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its
# template only branches on 'max', so the literal scan misses 'high'. Add it
# (matched on whole repo-name segments, so 'deepseek-v40' won't false-match)
# to expose the full none/high/max ladder instead of none/max.
segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1])
is_dsv4 = "deepseek4" in segments or any(
a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:])
)
if is_dsv4 and "high" not in effort_levels:
effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index)
# GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort
# level among a discrete set (e.g. 'high' | 'max'). Distinct from
# gpt-oss (reasoning_effort only, no on/off gate) and Qwen
@ -1741,9 +1751,13 @@ class LlamaCppBackend:
# 'low' effort the way gpt-oss does (those models genuinely
# cannot disable).
thinking_off = enable_thinking is False or reasoning_effort == "none"
if enable_thinking is not None or reasoning_effort == "none":
# A named effort level implies thinking on, so emit enable_thinking
# even if the caller sent only reasoning_effort (else the template
# defaults it off and the requested level never renders).
effort_on = reasoning_effort in self._reasoning_effort_levels
if enable_thinking is not None or reasoning_effort == "none" or effort_on:
kwargs["enable_thinking"] = not thinking_off
if not thinking_off and reasoning_effort in self._reasoning_effort_levels:
if not thinking_off and effort_on:
kwargs["reasoning_effort"] = reasoning_effort
elif self._reasoning_style == "reasoning_effort":
if reasoning_effort in ("none", "low", "medium", "high"):
@ -3129,6 +3143,12 @@ class LlamaCppBackend:
_CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch)
_CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x)
_CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok)
# DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large
# context-scaling compute buffer the rates above miss (present even with an f16
# cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without
# it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU.
_DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch
_DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M)
def _estimate_compute_buffer_bytes(
self,
@ -3178,6 +3198,14 @@ class LlamaCppBackend:
if n_embd <= 0 or n_ctx <= 0:
return 0
ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
if getattr(self, "_architecture", None) == "deepseek4":
# DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires
# for any KV type -- the indexer scratch is present even with an f16 cache.
ub_scale = ub / self._DEFAULT_N_UBATCH
return int(
self._DSV4_CTX_COMPUTE_FLAT_BYTES
+ self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale
)
if _kv_bytes_per_elem(cache_type_kv) < 2.0:
# Quantized cache: the dequant scratch dominates and scales with n_embd.
# MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on
@ -8575,13 +8603,31 @@ class LlamaCppBackend:
}
def _flush_reasoning_and_buffer():
"""Append buffered reasoning (as a <think> block) then the held
"""Close a live-streamed <think> block (or emit the buffered reasoning
as one block if it never streamed), then append the held
content_buffer to the cumulative display text."""
nonlocal cumulative_display
if reasoning_accum:
nonlocal cumulative_display, in_thinking
if in_thinking:
cumulative_display += "</think>"
in_thinking = False
elif reasoning_accum:
cumulative_display += "<think>" + reasoning_accum + "</think>"
cumulative_display += content_buffer
def _close_streamed_think() -> bool:
"""Close a live-streamed <think> before a tool call drains, so
consumers without a reasoning extractor (Anthropic) get a balanced
block. Returns True when the caller should yield the result."""
nonlocal cumulative_display, in_thinking, _last_emitted
if not in_thinking:
return False
cumulative_display += "</think>"
in_thinking = False
if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output:
_last_emitted = cumulative_display
return True
return False
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool:
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
probe = strip_llama3_leading_sentinels(text.lstrip())
@ -8769,6 +8815,10 @@ class LlamaCppBackend:
# the structured tool call.
has_structured_tc = True
detect_state = _S_DRAINING
# Close the reasoning prefix before the tool card
# (mirrors the is_match path).
if _close_streamed_think():
yield {"type": "content", "text": cumulative_display}
for tc_d in tc_deltas:
idx = tc_d.get("index", 0)
if idx not in tool_calls_acc:
@ -8854,17 +8904,17 @@ class LlamaCppBackend:
continue
# ── Reasoning tokens ──
# Yield only in STREAMING. In BUFFERING and
# DRAINING, accumulate silently so we don't
# corrupt the consumer's prev_text tracker
# (routes/inference.py never resets it
# between tool iterations).
# Stream live except while DRAINING: reasoning is
# orthogonal to tool detection (content_buffer
# only), and the route resets prev_text on
# tool_start, so the <think> block stays a
# monotonic prefix like the no-tool path.
reasoning = delta.get("reasoning_content", "")
if reasoning:
if _reasoning_started_at is None:
_reasoning_started_at = time.monotonic()
reasoning_accum += reasoning
if detect_state == _S_STREAMING:
if detect_state != _S_DRAINING:
if not in_thinking:
cumulative_display += "<think>"
in_thinking = True
@ -8992,9 +9042,15 @@ class LlamaCppBackend:
_hold_buffer = True
if _drain_silently:
# No visible prefix -- the buffered text IS
# the call; drain without yielding it.
# The buffered content IS the call; drain it
# without yielding. A live <think> prefix is
# separate from it -- close that.
detect_state = _S_DRAINING
if _close_streamed_think():
yield {
"type": "content",
"text": cumulative_display,
}
elif is_match:
# Tool signal -- flush any visible
# prefix before DRAINING so the
@ -9087,7 +9143,9 @@ class LlamaCppBackend:
),
}
elif reasoning_accum and not has_content_tokens:
# Reasoning-only reply: show it as plain text.
# Reasoning-only reply: show it as the main response,
# not a thinking block (mirrors the no-tool path; the
# route's extractor closes the streamed <think>).
if _reasoning_started_at is not None and not _reasoning_summary_emitted:
_reasoning_summary_emitted = True
yield _reasoning_summary_event(_reasoning_started_at)

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

@ -108,13 +108,11 @@ class InferenceOrchestrator:
@property
def default_models(self) -> list[str]:
# Wait up to 5s for background HF fetch
self._top_models_ready.wait(timeout = 5)
top_gguf = self._top_gguf_cache or []
top_hub = self._top_hub_cache or []
# Curated static defaults first, then HF download-ranked to backfill.
# Send extras so the frontend keeps 4 per category after removing
# downloaded ones.
# Never wait for the remote Hugging Face ranking during startup. Chat's
# first /api/models/list needs curated defaults immediately; the
# background fetch backfills extra choices on later calls.
result: list[str] = []
seen: set[str] = set()
for m in self._static_models + top_gguf + top_hub:

117
studio/backend/core/rag/embeddings.py
View file

@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None:
install_torchao_windows_rocm_stub()
class UnsafeEmbeddingModelError(RuntimeError):
"""Raised when the embedding model repo is flagged unsafe. A distinct type so the
llama-server fallback paths re-raise it instead of masking a security block as a
routine ST failure."""
def _ambient_hf_token() -> str | None:
"""The HF token the loader itself would use (HF_TOKEN env or the cached login), so
the scan can reach a gated/private repo instead of failing open. None if unavailable."""
try:
from huggingface_hub import get_token
return get_token()
except Exception:
return None
def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
"""The module directories a SentenceTransformer load reads weights from, taken from
the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``).
ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the
security scan: a flagged pickle directly under one must block. Returns () on any
failure (no modules.json, offline, malformed) so the guard never bricks the embedder.
"""
try:
import json
from utils.paths import is_local_path
if is_local_path(name):
from pathlib import Path
from utils.paths import normalize_path
path = Path(normalize_path(name)).expanduser() / "modules.json"
if not path.is_file():
return ()
data = json.loads(path.read_text())
else:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
try:
local = hf_hub_download(name, "modules.json", token = token or None)
except EntryNotFoundError:
return ()
data = json.loads(open(local).read())
subdirs = []
for module in data or ():
sub = str((module or {}).get("path", "")).strip().strip("/")
if sub:
subdirs.append(sub)
return tuple(dict.fromkeys(subdirs))
except Exception:
return ()
def _guard_model_security(name: str) -> None:
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
/settings gate (a name can also arrive via env/default); local paths and unreachable
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
"""
try:
from utils.security import evaluate_file_security, security_load_subdirs
token = _ambient_hf_token()
# Union the audio-model load roots with the ST module dirs so a flagged pickle
# directly under a Transformer module dir (0_Transformer/) blocks instead of
# passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
)
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
except Exception:
return
if blocked:
raise UnsafeEmbeddingModelError(
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
"scan; refusing to load. Set a different RAG embedding model."
)
def _get(model_name: str | None = None):
"""Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
for a ~1.5x speedup at negligible accuracy loss."""
@ -75,6 +156,7 @@ def _get(model_name: str | None = None):
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(
name, device = device, model_kwargs = {"torch_dtype": "float16"}
)
@ -159,6 +241,8 @@ class _SentenceTransformersBackend:
):
try:
return _st_encode(texts, model_name = model_name, normalize = normalize)
except UnsafeEmbeddingModelError:
raise # a security block must hard-fail, not fall back to llama-server
except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure
# ST loaded but this encode blew up; swap the process to the llama-server
# embedder (so later encodes stay in one space) and retry.
@ -222,6 +306,8 @@ def _build_st_backend_or_fallback():
try:
backend.warm(model_name = None)
return backend
except UnsafeEmbeddingModelError:
raise # a security block must hard-fail, not fall back to llama-server
except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure
fallback = _try_make_llama_backend()
if fallback is None:
@ -290,6 +376,37 @@ def _reset_backend() -> None:
_backend_key = None
def active_backend_is_llama() -> bool:
"""True when this process actually embeds via the llama-server (GGUF) backend.
Reflects the ACTUAL built backend once one exists: an ``auto`` install that
resolves to sentence-transformers but then falls back to llama-server at
runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or
``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so
callers gating on the ST pickle must see llama here. Before any backend is
built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw
key) exactly as a fresh process would. Never raises: a backend probe must not
block saving a model."""
try:
with _backend_lock:
backend = _backend
if backend is not None:
# A backend exists: report what it ACTUALLY is. A concrete
# sentence-transformers backend must return False even if the
# resolver would now pick llama, so its pickle stays gated. If the
# llama import fails we cannot be llama, so fall to the safe False.
try:
from .embed_llama_server import LlamaServerBackend
except Exception: # noqa: BLE001 - llama plumbing import must never block
return False
return isinstance(backend, LlamaServerBackend)
raw = (config.EMBED_BACKEND or "auto").strip().lower()
key = _resolve_auto() if raw in _AUTO_ALIASES else raw
return key in _LLAMA_ALIASES
except Exception: # noqa: BLE001 - a backend probe must never block saving
return False
def warm(model_name: str | None = None) -> None:
"""Eagerly load the embedder so the first real request isn't slow."""
_get_backend().warm(model_name = model_name)

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

@ -56,6 +56,7 @@ class ChatThread(BaseModel):
projectId: Optional[str] = None
archived: bool = False
createdAt: int
updatedAt: Optional[int] = None
openaiCodeExecContainerId: Optional[str] = None
anthropicCodeExecContainerId: Optional[str] = None
forkedFromThreadId: Optional[str] = None
@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel):
projectId: Optional[str] = None
archived: Optional[bool] = None
createdAt: Optional[int] = None
updatedAt: Optional[int] = None
openaiCodeExecContainerId: Optional[str] = None
anthropicCodeExecContainerId: Optional[str] = None
@ -252,7 +254,7 @@ async def patch_thread(
current_subject: str = Depends(get_current_subject),
):
patch = payload.model_dump(exclude_unset = True)
for field in ("title", "modelType", "modelId", "archived", "createdAt"):
for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"):
if field in patch and patch[field] is None:
raise HTTPException(status_code = 400, detail = f"{field} cannot be null")
if patch.get("projectId") and get_chat_project(patch["projectId"]) is None:

78
studio/backend/routes/settings.py
View file

@ -260,17 +260,29 @@ def _embedding_model_response() -> EmbeddingModelResponse:
)
def _llama_backend_active() -> bool:
"""True when this install embeds via the llama-server (GGUF) backend."""
from core.rag import config as rag_config
from core.rag import embeddings
def _ambient_hf_token() -> Optional[str]:
"""The HF token the loader would use (HF_TOKEN env or the cached login), so a gated
repo is scanned rather than failing open. None if unavailable."""
try:
raw = (rag_config.EMBED_BACKEND or "auto").strip().lower()
key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw
from huggingface_hub import get_token
return get_token()
except Exception:
return None
def _llama_backend_active() -> bool:
"""True when this install actually embeds via the llama-server (GGUF) backend.
Delegates to the embeddings module so a runtime fallback from
sentence-transformers to llama-server (after a torch/CUDA load or encode
failure) is honored: in that state the process loads only inert GGUF, so the
ST pickle gate below must not hard-block a repo whose GGUF companion is clean.
Before any backend is built this still reflects the resolver."""
from core.rag import embeddings
try:
return embeddings.active_backend_is_llama()
except Exception: # noqa: BLE001 - backend probe must never block saving
return False
return key in embeddings._LLAMA_ALIASES
def _resolves_as_local_gguf(model: str) -> bool:
@ -357,6 +369,8 @@ def update_embedding_model(
"""Set the RAG embedding model. Unless ``force`` is set, the repo is verified
to be an embedding model via HF metadata; an unverifiable model (wrong type,
typo, gated repo, or no network) returns 409 so the UI can offer "save anyway".
A repo flagged unsafe by HF's security scan returns 403 instead: a hard block
that ``force`` cannot bypass, so the UI must not offer "save anyway".
Documents indexed under the previous model must be re-uploaded."""
from utils.models import is_embedding_model
@ -370,15 +384,51 @@ def update_embedding_model(
event = "settings.update_embedding_model_failed",
log = logger,
) from exc
hf_token = (payload.hf_token or "").strip() or None
# The env/default model needs no verification; saving it is a no-op override.
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
# what the backend loads, and HF metadata cannot verify a local path.
if (
model != default_embedding_model()
and not payload.force
and not (_llama_backend_active() and _resolves_as_local_gguf(model))
):
hf_token = (payload.hf_token or "").strip() or None
is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model)
# The pickle gate only matters for the sentence-transformers backend, which is what
# deserializes pickles. On the llama-server backend the embedder loads GGUF files
# (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would
# wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability
# checks below cover that path instead.
scan_st_pickle = (
model != default_embedding_model() and not is_local_gguf and not _llama_backend_active()
)
if scan_st_pickle:
# Malware/pickle gate before we persist a repo the embedder later loads with
# SentenceTransformer. Runs even under force (force only skips the is-embedding
# type check for offline/local repos HF cannot verify); local paths and
# unreachable scans fail open inside evaluate_file_security.
from utils.security import evaluate_file_security, security_load_subdirs
from core.rag.embeddings import _st_module_subdirs
# Fall back to the loader's own token so a gated/private repo is actually scanned
# (a token-less scan fails open for exactly the repo that would still load).
scan_token = hf_token or _ambient_hf_token()
# Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
# one blocks instead of passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys(
(
*security_load_subdirs(model, scan_token),
*_st_module_subdirs(model, scan_token),
)
)
)
if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
# 403, not 409: the client routes every 409 into the forceable "save anyway"
# flow, but this block is a hard, non-forceable security refusal.
raise HTTPException(
status_code = 403,
detail = (
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
"cannot be used as the embedding model."
),
)
if model != default_embedding_model() and not payload.force and not is_local_gguf:
from core.rag import config as rag_config
# A GGUF-named repo on the llama-server backend is loaded from its .gguf

79
studio/backend/run.py
View file

@ -12,6 +12,79 @@ import time
from pathlib import Path
from typing import Optional
def _fix_torch_cuda_ld_path():
"""Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH.
PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, ...) in
``site-packages/nvidia/*/lib``. On Linux the dynamic linker reads
LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a
pre-existing LD_LIBRARY_PATH pointing at a different system CUDA (e.g.
/usr/local/cuda-13/lib64 from conda or a Docker base image) shadows torch's
libs and triggers "undefined symbol" errors when torch is imported. Detect
torch's lib dirs (without importing torch) and prepend them. Returns True if
LD_LIBRARY_PATH was changed.
"""
if sys.platform != "linux":
return False
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
if not ld_path:
return False
try:
import importlib.util
spec = importlib.util.find_spec("torch")
if not spec or not spec.origin:
return False
torch_dir = os.path.dirname(spec.origin)
site_pkgs = os.path.dirname(torch_dir)
nvidia_dir = os.path.join(site_pkgs, "nvidia")
lib_dirs = []
torch_lib = os.path.join(torch_dir, "lib")
if os.path.isdir(torch_lib):
lib_dirs.append(torch_lib)
if os.path.isdir(nvidia_dir):
for sub in sorted(os.listdir(nvidia_dir)):
lib = os.path.join(nvidia_dir, sub, "lib")
if os.path.isdir(lib):
lib_dirs.append(lib)
if not lib_dirs:
return False
existing = ld_path.split(":")
if existing[: len(lib_dirs)] == lib_dirs:
return False # already at the front, nothing to do
torch_set = set(lib_dirs)
cleaned = [p for p in existing if p not in torch_set]
os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned)
return True
except Exception:
return False
_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED"
def _maybe_reexec_for_cuda_ld_path():
"""Re-exec once so the dynamic linker sees the corrected LD_LIBRARY_PATH.
LD_LIBRARY_PATH is read at process start, so editing os.environ in-process
cannot fix the running interpreter; a single re-exec is required. Call only
from a true entry point (the ``if __name__ == "__main__"`` block), never at
import time, because os.execv replaces the whole process (an embedder such
as Colab that does ``from run import run_server`` must not be re-exec'd).
"""
if _LD_FIXED_SENTINEL in os.environ:
return
if not _fix_torch_cuda_ld_path():
return
os.environ[_LD_FIXED_SENTINEL] = "1"
argv = getattr(sys, "orig_argv", None) or [sys.executable, *sys.argv]
os.execv(sys.executable, argv)
# Suppress C-level dependency warnings globally (e.g. SwigPyPacked).
os.environ["PYTHONWARNINGS"] = "ignore"
@ -1457,6 +1530,12 @@ def _build_arg_parser():
# For direct execution (also invoked by CLI via os.execvp / subprocess).
if __name__ == "__main__":
# Correct a conflicting system CUDA on LD_LIBRARY_PATH before torch is
# imported (below, via run_server). Re-execs once on Linux so the dynamic
# linker uses torch's bundled CUDA libs; no-op on other platforms, when
# LD_LIBRARY_PATH is unset or already correct, or after the single re-exec.
_maybe_reexec_for_cuda_ld_path()
import signal
import traceback

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

@ -240,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
project_id TEXT,
archived INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER,
openai_code_exec_container_id TEXT,
anthropic_code_exec_container_id TEXT,
forked_from_thread_id TEXT,
@ -261,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT")
if "forked_from_message_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT")
if "updated_at" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER")
# Floor at created_at: forked threads copy older ancestor messages,
# so the fork's creation time must win over the branch message times.
conn.execute(
"""
UPDATE chat_threads SET updated_at = MAX(
COALESCE(
(
SELECT MAX(m.created_at) FROM chat_messages m
WHERE m.thread_id = chat_threads.id
),
created_at
),
created_at
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
@ -992,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict:
"projectId": data.get("project_id") or None,
"archived": bool(data["archived"]),
"createdAt": data["created_at"],
"updatedAt": data.get("updated_at")
if data.get("updated_at") is not None
else data["created_at"],
"openaiCodeExecContainerId": data.get("openai_code_exec_container_id"),
"anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"),
"forkedFromThreadId": data.get("forked_from_thread_id"),
@ -1039,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict:
conn.execute(
"""
INSERT INTO chat_threads
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
model_type = excluded.model_type,
@ -1049,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict:
project_id = excluded.project_id,
archived = excluded.archived,
created_at = excluded.created_at,
updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at),
openai_code_exec_container_id = excluded.openai_code_exec_container_id,
anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id,
forked_from_thread_id = excluded.forked_from_thread_id,
@ -1063,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict:
thread.get("projectId"),
1 if thread.get("archived") else 0,
int(thread["createdAt"]),
int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None,
thread.get("openaiCodeExecContainerId"),
thread.get("anthropicCodeExecContainerId"),
thread.get("forkedFromThreadId"),
@ -1084,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]:
"projectId": ("project_id", patch.get("projectId")),
"archived": ("archived", 1 if patch.get("archived") else 0),
"createdAt": ("created_at", patch.get("createdAt")),
"updatedAt": ("updated_at", patch.get("updatedAt")),
"openaiCodeExecContainerId": (
"openai_code_exec_container_id",
patch.get("openaiCodeExecContainerId"),
@ -1155,7 +1180,8 @@ def list_chat_threads(
conn = get_connection()
try:
rows = conn.execute(
f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC",
f"SELECT * FROM chat_threads {where} "
"ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC",
values,
).fetchall()
return [_chat_thread_from_row(row) for row in rows]
@ -1394,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts(
)
def _bump_chat_thread_updated_at(
conn: sqlite3.Connection, thread_id: str, message_created_at: int
) -> None:
conn.execute(
"""
UPDATE chat_threads
SET updated_at = MAX(COALESCE(updated_at, created_at), ?)
WHERE id = ?
""",
(message_created_at, thread_id),
)
def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None:
"""Set updated_at from the remaining messages, floored at created_at.
Unlike the ratchet-only bump, this can lower updated_at -- needed after
pruning, which may delete the thread's newest message.
"""
conn.execute(
"""
UPDATE chat_threads
SET updated_at = MAX(
COALESCE(
(
SELECT MAX(m.created_at) FROM chat_messages m
WHERE m.thread_id = chat_threads.id
),
created_at
),
created_at
)
WHERE id = ?
""",
(thread_id,),
)
def upsert_chat_message(message: dict) -> dict:
conn = get_connection()
try:
@ -1432,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict:
int(message["createdAt"]),
),
)
_bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"]))
conn.commit()
return message
except Exception:
@ -1484,6 +1549,12 @@ def sync_chat_messages(
for m in messages
],
)
if prune_missing:
_recompute_chat_thread_updated_at(conn, thread_id)
elif messages:
_bump_chat_thread_updated_at(
conn, thread_id, max(int(m["createdAt"]) for m in messages)
)
conn.commit()
return list_chat_messages(thread_id)
except ChatMessageConflictError:

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

@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO
from types import SimpleNamespace
def _emitter_client_text(events: list[str]) -> str:
"""Concatenate the text_delta payloads an SSE event list carries."""
text = ""
for line in events:
for raw in line.split("\n"):
raw = raw.strip()
if not raw.startswith("data: "):
continue
data = json.loads(raw[len("data: ") :])
delta = data.get("delta", {})
if delta.get("type") == "text_delta":
text += delta.get("text", "")
return text
def test_anthropic_emitter_closes_reasoning_only_think_block():
# A reasoning-only reply streams <think>X live then shrinks to bare X at EOF.
# This emitter diffs cumulative snapshots and drops the shrink, so without a
# closing pass the client text would end on an unclosed <think>. finish()
# must balance it.
emitter = AnthropicStreamEmitter()
events = emitter.start("msg_1", "m")
events += emitter.feed({"type": "content", "text": "<think>The capital"})
events += emitter.feed({"type": "content", "text": "<think>The capital of France is Paris."})
# The generator's final bare-text shrink (dropped by the cumulative diff).
events += emitter.feed({"type": "content", "text": "The capital of France is Paris."})
events += emitter.finish()
assert _emitter_client_text(events) == "<think>The capital of France is Paris.</think>"
def test_anthropic_emitter_does_not_double_close_balanced_think():
# A reasoning-then-answer reply already closes its own </think>; the balancer
# must not append a second one.
emitter = AnthropicStreamEmitter()
events = emitter.start("msg_1", "m")
events += emitter.feed({"type": "content", "text": "<think>Thinking."})
events += emitter.feed({"type": "content", "text": "<think>Thinking.</think>Answer."})
events += emitter.finish()
assert _emitter_client_text(events) == "<think>Thinking.</think>Answer."
def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch):
import routes.inference as inf_mod

134
studio/backend/tests/test_chat_history_storage.py
View file

@ -4,6 +4,7 @@
import os
import platform
import shutil
import sqlite3
import threading
import uuid
from pathlib import Path
@ -11,6 +12,7 @@ from pathlib import Path
import pytest
from storage import studio_db
from utils.paths import studio_db_path
def _reset_studio_db(
@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch):
assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}]
def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
thread = studio_db.upsert_chat_thread(_thread())
assert thread["updatedAt"] == thread["createdAt"]
studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi"))
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old"))
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
studio_db.sync_chat_messages(
"thread-1",
[_message("msg-2", 1_700_000_001_000, "newer")],
)
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000
def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
thread = studio_db.upsert_chat_thread(_thread())
studio_db.sync_chat_messages(
"thread-1",
[
_message("msg-1", 1_700_000_000_500, "older"),
_message("msg-2", 1_700_000_001_000, "newest"),
],
prune_missing = True,
)
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000
# Pruning the newest message must lower updated_at to the remaining one.
studio_db.sync_chat_messages(
"thread-1",
[_message("msg-1", 1_700_000_000_500, "older")],
prune_missing = True,
)
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
# Pruning every message falls back to created_at.
studio_db.sync_chat_messages("thread-1", [], prune_missing = True)
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"]
def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi"))
studio_db.upsert_chat_thread(_thread())
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
older = _thread("thread-old")
older["createdAt"] = 1_700_000_000_000
newer = _thread("thread-new")
newer["createdAt"] = 1_700_000_100_000
studio_db.upsert_chat_thread(older)
studio_db.upsert_chat_thread(newer)
assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"]
studio_db.upsert_chat_message(
_message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old")
)
assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"]
def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
db_path = studio_db_path()
db_path.parent.mkdir(parents = True, exist_ok = True)
conn = sqlite3.connect(str(db_path))
try:
conn.execute(
"""
CREATE TABLE chat_threads (
id TEXT NOT NULL PRIMARY KEY,
title TEXT NOT NULL,
model_type TEXT NOT NULL,
model_id TEXT,
pair_id TEXT,
archived INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE chat_messages (
id TEXT NOT NULL PRIMARY KEY,
thread_id TEXT NOT NULL,
parent_id TEXT,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
attachments_json TEXT,
metadata_json TEXT,
created_at INTEGER NOT NULL
)
"""
)
conn.execute(
"INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)",
("thread-with-msgs", "Old", "base", 1_700_000_000_000),
)
conn.execute(
"INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)",
("thread-empty", "Empty", "base", 1_700_000_050_000),
)
# Fork-like thread: copied ancestor messages predate the thread itself.
conn.execute(
"INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)",
("thread-fork", "Fork", "base", 1_700_000_100_000),
)
conn.executemany(
"INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)",
[
("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000),
("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000),
("m3", "thread-fork", "user", "[]", 1_700_000_001_000),
],
)
conn.commit()
finally:
conn.close()
assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000
assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000
assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000
def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
project = studio_db.upsert_chat_project(_project())

61
studio/backend/tests/test_compute_buffer.py
View file

@ -65,12 +65,14 @@ def _backend(
vocab = 248320,
embd = 5120,
mla = None,
arch = None,
):
"""Backend with just the dims the compute-buffer estimate reads."""
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._vocab_size = vocab
b._embedding_length = embd
b._key_length_mla = mla # non-None -> MLA (compressed attention)
b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4')
return b
@ -290,3 +292,62 @@ class TestContextBufferMLA:
b = _backend(embd = 6144, mla = 256)
est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB
assert est <= 4141 * 1.7
class TestContextBufferDSV4:
"""DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention
compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache).
Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit
must see this so it does not commit the full 1M train context and OOM (spilling
to CPU at ~4 tok/s)."""
_MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx
GIB = 1024**3
def test_covers_measured_1m_buffer(self):
b = _backend(embd = 4096, arch = "deepseek4")
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}"
def test_not_wildly_over_at_1m(self):
# Within ~1.3x of measured so the fit still grants a large (~256k) context.
b = _backend(embd = 4096, arch = "deepseek4")
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
assert gib <= self._MEASURED_1M_GIB * 1.3
def test_fires_for_f16_cache(self):
# The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must
# reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx.
dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes(
262144, cache_type_kv = "f16"
)
other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes(
262144, cache_type_kv = "f16"
)
assert dsv4 > 40 * other
def test_cache_type_independent(self):
# Indexer scratch is present for an f16 and a quantized cache alike.
b = _backend(embd = 4096, arch = "deepseek4")
assert b._compute_buffer_ctx_bytes(
262144, cache_type_kv = "f16"
) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
def test_flat_floor_at_small_ctx(self):
# ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB).
b = _backend(embd = 4096, arch = "deepseek4")
assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0
def test_scales_with_context_and_ubatch(self):
b = _backend(embd = 4096, arch = "deepseek4")
assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536)
assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes(
131072, n_ubatch = 256
)
def test_non_dsv4_unchanged(self):
# Regression guard: a non-deepseek4 model keeps the mask-only f16 rate.
b = _backend(embd = 4096, arch = "llama")
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000
expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY
assert per_tok == pytest.approx(expected, rel = 1e-6)

181
studio/backend/tests/test_deepseek_v4_thinking_effort.py Normal file
View file

@ -0,0 +1,181 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""DeepSeek-V4-Flash reasoning toggle: None / High / Max.
The GGUF template gates thinking with ``enable_thinking`` and only branches
``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking).
Detection used to return the single level ``['max']``, so the UI collapsed to
None / Max and the plain-thinking tier was unreachable. Detection now surfaces
``'high'`` as that plain tier, giving None / High / Max. These tests pin the
classifier, the GLM-style parity case, and the full request-kwargs -> rendered
prompt path for each state (the model itself is too large to load here).
"""
from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
_backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking
# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think
# fallback. Any non-'max' effort renders as ordinary thinking.
DEEPSEEK_V4_TEMPLATE = """
{%- if not thinking is defined -%}
{%- if enable_thinking is defined -%}
{%- set thinking = enable_thinking -%}
{%- else -%}
{%- set thinking = false -%}
{%- endif -%}
{%- endif -%}
{%- if not reasoning_effort is defined -%}
{%- set reasoning_effort = none -%}
{%- endif -%}
{{- bos_token -}}
{%- if thinking and reasoning_effort == 'max' -%}
{{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}}
{%- endif -%}
{%- for message in messages -%}
{{- '<|User|>' + (message['content'] or '') -}}
{%- endfor -%}
{%- if add_generation_prompt -%}
{{- '<|Assistant|>' -}}
{%- if thinking -%}{{- '<think>' -}}{%- else -%}{{- '</think>' -}}{%- endif -%}
{%- endif -%}
"""
# GLM-5.2-style: branches on two effort literals, so 'high' already exists as
# the sub-'max' tier and detection must leave the pair untouched.
GLM_STYLE_TEMPLATE = """
{%- if enable_thinking -%}
{%- if reasoning_effort == 'high' -%}{{- 'H' -}}
{%- elif reasoning_effort == 'max' -%}{{- 'M' -}}
{%- endif -%}
{%- endif -%}
"""
# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped
# to deepseek-v4, so this must stay ['max'] (no phantom 'high').
NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE
# A template whose sole effort literal is a sub-'max' level: the guard targets
# only the ['max']-alone case, so a lone 'high' stays a singleton.
HIGH_ONLY_TEMPLATE = """
{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%}
"""
def _render(template: str, **kwargs) -> str:
jinja2 = pytest.importorskip("jinja2")
env = jinja2.Environment()
tmpl = env.from_string(template)
return tmpl.render(bos_token = "<BOS>", add_generation_prompt = True, **kwargs)
# -- Classifier -------------------------------------------------------
def test_deepseek_v4_surfaces_high_as_plain_tier():
"""Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash")
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
def test_glm_style_two_level_template_unchanged():
"""A template that already names a sub-'max' tier is left as-is."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2")
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
def test_synthetic_high_scoped_to_deepseek_v4():
"""The same ['max']-only template under a non-deepseek id keeps ['max']."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF")
assert flags["reasoning_effort_levels"] == ["max"]
def test_guard_does_not_fire_for_sub_max_singleton():
"""The expansion targets only ['max']; a lone 'high' stays a singleton."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only")
assert flags["reasoning_effort_levels"] == ["high"]
# -- Request kwargs -> rendered prompt, for each state ----------------
def _kwargs_for(flags: dict, enable_thinking, reasoning_effort):
"""Drive the real backend method with a shim carrying the detected flags."""
from core.inference.llama_cpp import LlamaCppBackend
shim = SimpleNamespace(
_supports_reasoning = flags["supports_reasoning"],
_reasoning_always_on = flags["reasoning_always_on"],
_reasoning_style = flags["reasoning_style"],
_reasoning_effort_levels = flags["reasoning_effort_levels"],
_supports_preserve_thinking = flags["supports_preserve_thinking"],
)
build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim)
return build(enable_thinking, reasoning_effort, None) or {}
def _flags():
from core.inference.llama_cpp import detect_reasoning_flags
return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash")
def test_none_state_renders_non_thinking():
"""UI 'None' -> enable_thinking=false -> closed </think>, no preamble."""
kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None)
assert kwargs == {"enable_thinking": False}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("</think>")
assert "Absolute maximum" not in out
def test_high_state_renders_plain_thinking():
"""UI 'High' -> et=true, effort=high -> open <think>, no max preamble."""
kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high")
assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("<think>")
assert "Absolute maximum" not in out
def test_max_state_injects_max_preamble():
"""UI 'Max' -> et=true, effort=max -> open <think> plus the max preamble."""
kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max")
assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("<think>")
assert "Absolute maximum" in out
def test_high_effort_alone_enables_thinking():
"""API caller sending only reasoning_effort='high' (no enable_thinking) still
gets thinking on, so the newly exposed High mode renders correctly."""
kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high")
assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("<think>")
assert "Absolute maximum" not in out

365
studio/backend/tests/test_embedding_model_security_gate.py Normal file
View file

@ -0,0 +1,365 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The RAG embedding model must pass the malware/pickle gate before it is persisted or
loaded. A flagged repo (or any repo saved with force) previously reached
SentenceTransformer unscanned, bypassing the normal model-load protections."""
from pathlib import Path
import sys
import types as _types
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
import routes.settings as settings
class _Decision:
def __init__(self, blocked):
self.blocked = blocked
def _security_stub(blocked):
mod = _types.ModuleType("utils.security")
mod.evaluate_file_security = lambda *a, **k: _Decision(blocked)
mod.security_load_subdirs = lambda *a, **k: ()
return mod
@pytest.fixture
def client(monkeypatch):
# The settings scan unions in the ST module dirs read from modules.json; keep it
# offline and deterministic for the endpoint tests that use this fixture.
import core.rag.embeddings as embeddings
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
return TestClient(app, raise_server_exceptions = False), saved
def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch):
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
r = c.put(
"/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True}
)
# 403, not the forceable 409, so the client does not offer "save anyway".
assert r.status_code == 403
assert "model" not in saved # force must not persist a flagged repo
def test_flagged_repo_is_blocked_without_force(client, monkeypatch):
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
assert r.status_code == 403
assert "model" not in saved
def test_hard_block_uses_non_forceable_status(client, monkeypatch):
# The forceable verification path uses 409; the hard security block must be distinct
# (403) so the frontend never routes it into the "save anyway" force flow.
c, _saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
assert blocked.status_code == 403
# A verification failure (not-an-embedding-model) stays forceable at 409.
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False)
import utils.models as _models
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"})
assert unverified.status_code == 409
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
# force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped.
called = {"scanned": False}
mod = _types.ModuleType("utils.security")
def _fail(*a, **k):
called["scanned"] = True
return _Decision(True)
mod.evaluate_file_security = _fail
mod.security_load_subdirs = lambda *a, **k: ()
monkeypatch.setitem(sys.modules, "utils.security", mod)
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
c = TestClient(app, raise_server_exceptions = False)
r = c.put(
"/embedding-model",
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
)
assert r.status_code == 200
assert called["scanned"] is False # the ST pickle scan never ran on the llama path
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch):
# auto resolves to sentence-transformers (GPU present) but the embedder fell back to
# llama-server at runtime (torch/CUDA load or encode failure), so the process now loads
# only inert GGUF. The real _llama_backend_active() must reflect that cached fallback,
# so a flagged ST repo with a clean GGUF companion must not be hard-blocked here.
import core.rag.embeddings as embeddings
from core.rag.embed_llama_server import LlamaServerBackend
# Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even
# though the auto resolver would still say sentence-transformers.
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
# Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the
# real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored.
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
called = {"scanned": False}
mod = _types.ModuleType("utils.security")
def _fail(*a, **k):
called["scanned"] = True
return _Decision(True)
mod.evaluate_file_security = _fail
mod.security_load_subdirs = lambda *a, **k: ()
monkeypatch.setitem(sys.modules, "utils.security", mod)
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
c = TestClient(app, raise_server_exceptions = False)
r = c.put(
"/embedding-model",
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
)
assert r.status_code == 200
assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch):
# active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers
# to the resolver (fresh-process behavior) when none has been built yet.
import core.rag.embeddings as embeddings
import core.rag.config as rag_config
from core.rag.embed_llama_server import LlamaServerBackend
# A cached llama backend wins even when auto would resolve to sentence-transformers.
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto")
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
assert embeddings.active_backend_is_llama() is True
# A cached ST backend reports False even when the resolver now picks llama, so its
# pickle stays gated (the cached backend, not the resolver, is what actually embeds).
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend())
assert embeddings.active_backend_is_llama() is False
# No cached backend -> the resolver decides, unchanged from before.
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
monkeypatch.setattr(embeddings, "_backend", None)
assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
assert embeddings.active_backend_is_llama() is True # auto -> llama-server
# An explicit (non-auto) key is honored verbatim without a cached backend.
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server")
assert embeddings.active_backend_is_llama() is True
def test_settings_scan_scopes_module_subdirs(monkeypatch):
# The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a
# pickle directly under one blocks; assert those subdirs reach evaluate_file_security.
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
import core.rag.embeddings as embeddings
monkeypatch.setattr(
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
)
seen = {}
def _capture(*a, **k):
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
return _Decision(False)
mod = _types.ModuleType("utils.security")
mod.security_load_subdirs = lambda *a, **k: ()
mod.evaluate_file_security = _capture
monkeypatch.setitem(sys.modules, "utils.security", mod)
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
c = TestClient(app, raise_server_exceptions = False)
r = c.put(
"/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True}
)
assert r.status_code == 200
assert "0_Transformer" in seen["subdirs"]
def test_clean_repo_saves_under_force(client, monkeypatch):
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True})
assert r.status_code == 200
assert saved.get("model") == "acme/clean-embed"
def test_load_sink_refuses_flagged_model(monkeypatch):
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
import core.rag.embeddings as embeddings
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
embeddings._guard_model_security("attacker/malicious-embed")
def test_load_sink_allows_clean_model(monkeypatch):
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
import core.rag.embeddings as embeddings
embeddings._guard_model_security("acme/clean-embed") # no raise
def test_sink_threads_ambient_token_into_scan(monkeypatch):
# A gated repo set via env/default has no request token; the guard must feed the
# loader's own token to the scan, or it fails open for the repo that still loads.
seen = {}
mod = _types.ModuleType("utils.security")
mod.security_load_subdirs = (
lambda name, token = None: seen.setdefault("subdirs_token", token) or ()
)
mod.evaluate_file_security = lambda *a, **k: seen.setdefault(
"scan_token", k.get("hf_token")
) or _Decision(False)
monkeypatch.setitem(sys.modules, "utils.security", mod)
import core.rag.embeddings as embeddings
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient")
embeddings._guard_model_security("acme/gated-embed")
assert seen["scan_token"] == "hf_ambient"
assert seen["subdirs_token"] == "hf_ambient"
def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch):
# A flagged pickle directly under a Transformer module dir (0_Transformer/) must
# reach the scan as a load root; assert the guard unions the module dirs into
# load_subdirs so evaluate_file_security treats such a pickle as root-level.
seen = {}
def _capture(*a, **k):
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
return _Decision(False)
mod = _types.ModuleType("utils.security")
mod.security_load_subdirs = lambda name, token = None: ()
mod.evaluate_file_security = _capture
monkeypatch.setitem(sys.modules, "utils.security", mod)
import core.rag.embeddings as embeddings
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None)
monkeypatch.setattr(
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
)
embeddings._guard_model_security("acme/embed-with-module-dir")
assert "0_Transformer" in seen["subdirs"]
def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch):
# The helper must parse each module's non-empty "path" from a local repo's
# modules.json and drop the root-level ("") Transformer entry.
import json
import core.rag.embeddings as embeddings
(tmp_path / "modules.json").write_text(
json.dumps(
[
{"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."},
{"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."},
{"idx": 2, "name": "2", "path": "", "type": "..."},
]
)
)
subdirs = embeddings._st_module_subdirs(str(tmp_path), None)
assert subdirs == ("0_Transformer", "1_Pooling")
def test_st_module_subdirs_swallows_errors(monkeypatch):
# Any failure (no modules.json, offline, malformed) returns () so the guard never
# bricks the embedder.
import huggingface_hub
import core.rag.embeddings as embeddings
def _boom(*a, **k):
raise RuntimeError("offline")
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom)
assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == ()
def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch):
# The ST encode fallback must re-raise a security block, not swap to llama-server.
import core.rag.embeddings as embeddings
def _boom(*a, **k):
raise embeddings.UnsafeEmbeddingModelError("flagged")
monkeypatch.setattr(embeddings, "_st_encode", _boom)
monkeypatch.setattr(
embeddings,
"_switch_to_llama_fallback",
lambda err: pytest.fail("security block must not fall back to llama-server"),
)
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
embeddings._SentenceTransformersBackend().encode(["hi"])

42
studio/backend/tests/test_inference_default_models_non_blocking.py Normal file
View file

@ -0,0 +1,42 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Default Chat model metadata must not block on remote Hugging Face discovery."""
from __future__ import annotations
import sys
import time
from pathlib import Path
_BACKEND = Path(__file__).resolve().parent.parent
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from core.inference.orchestrator import InferenceOrchestrator # noqa: E402
def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch):
sleep_seconds = 2.0
def _slow_fetch(self: InferenceOrchestrator) -> None:
time.sleep(sleep_seconds)
self._top_gguf_cache = ["unsloth/slow-GGUF"]
self._top_models_ready.set()
monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch)
orchestrator = InferenceOrchestrator()
started = time.monotonic()
defaults = orchestrator.default_models
elapsed = time.monotonic() - started
assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s"
assert defaults == orchestrator._static_models
assert "unsloth/slow-GGUF" not in defaults
deadline = time.monotonic() + sleep_seconds + 5
while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline:
time.sleep(0.05)
assert "unsloth/slow-GGUF" in orchestrator.default_models

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

@ -221,7 +221,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html"
def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch):
def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch):
stream = [
_sse({"reasoning_content": "I am thinking."}),
_sse({"reasoning_content": " Still thinking."}),
@ -240,17 +240,236 @@ def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch):
)
)
content_texts = [e["text"] for e in events if e["type"] == "content"]
# Reasoning streams live during BUFFERING instead of arriving as one block:
# each reasoning delta is emitted immediately, wrapped in <think>.
assert content_texts[0] == "<think>I am thinking."
assert content_texts[1] == "<think>I am thinking. Still thinking."
# The final event closes the block and appends the answer.
assert content_texts[-1] == "<think>I am thinking. Still thinking.</think>Final answer."
summary_index = next(
i for i, event in enumerate(events) if event["type"] == "reasoning_summary"
)
content_index = next(i for i, event in enumerate(events) if event["type"] == "content")
assert summary_index < content_index
final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content")
assert summary_index < final_content_index
assert events[summary_index]["duration_ms"] == 62000
assert (
events[content_index]["text"]
== "<think>I am thinking. Still thinking.</think>Final answer."
def test_reasoning_streams_incrementally_with_tools(monkeypatch):
# Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the
# tool-loop generator must stream reasoning token-by-token like the no-tool
# path, not accumulate it and dump one buffered <think> block.
stream = [
_sse({"reasoning_content": "Step one."}),
_sse({"reasoning_content": " Step two."}),
_sse({"reasoning_content": " Step three."}),
_sse({"content": "Done."}),
_done(),
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream], payloads)
_patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0])
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "think then answer"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
reasoning_stage = [
e["text"]
for e in events
if e["type"] == "content"
and e["text"].startswith("<think>")
and "</think>" not in e["text"]
]
# One live emission per reasoning delta -- not a single dump.
assert reasoning_stage == [
"<think>Step one.",
"<think>Step one. Step two.",
"<think>Step one. Step two. Step three.",
]
final = [e["text"] for e in events if e["type"] == "content"][-1]
assert final == "<think>Step one. Step two. Step three.</think>Done."
def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch):
# A reasoning-only turn (whole answer in reasoning_content, no content, no
# tool) with a tool active streams the reasoning live, then resolves to the
# bare reasoning text -- identical to the no-tool generate_chat_completion
# path -- so the non-streaming drain still returns it as `content`, not an
# empty answer.
stream = [
_sse({"reasoning_content": "The capital of France is Paris."}),
_done(),
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream], payloads)
_patch_monotonic(monkeypatch, [1.0, 5.0, 5.0])
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "just think"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
content_texts = [e["text"] for e in events if e["type"] == "content"]
# Reasoning streamed live during BUFFERING (the fix).
assert content_texts[0] == "<think>The capital of France is Paris."
# Resolves to bare reasoning, matching the no-tool sibling.
assert content_texts[-1] == "The capital of France is Paris."
def test_reasoning_before_structured_tool_closes_think_block(monkeypatch):
# Regression: reasoning streamed live during BUFFERING must be closed with
# </think> before a structured tool_call drains, so consumers without a
# reasoning extractor (Anthropic /v1/messages) never receive an unclosed
# <think>. Mirrors the is_match (XML tool signal) path.
tool_stream = [
_sse({"reasoning_content": "Let me search."}),
*_structured_tool_call("web_search", {"query": "weather"}, "call_1"),
]
final_stream = [
_sse({"content": "It is sunny."}),
_done(),
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads)
_patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0])
monkeypatch.setattr(
"core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny"
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "weather?"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start")
content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"]
# Reasoning streamed live, then closed before the tool -- balanced block.
assert content_before_tool[0] == "<think>Let me search."
assert content_before_tool[-1] == "<think>Let me search.</think>"
def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]:
"""Replay the route's cumulative suffix-diff + reasoning extractor (the
shared core of routes/inference.py gguf_stream_chunks and the tool-loop
consumer) over content snapshots. Returns (visible, reasoning)."""
from routes.inference import _ResponsesReasoningExtractor
extractor = _ResponsesReasoningExtractor(parse_think_markers = True)
prev_text = ""
visible: list[str] = []
reasoning: list[str] = []
for cumulative in cumulatives:
new_text = cumulative[len(prev_text) :]
prev_text = cumulative
if not new_text:
continue
reasoning_delta, visible_delta = extractor.feed(new_text)
if reasoning_delta:
reasoning.append(reasoning_delta)
if visible_delta:
visible.append(visible_delta)
final_reasoning, final_visible = extractor.finish()
if final_reasoning:
reasoning.append(final_reasoning)
if final_visible:
visible.append(final_visible)
return "".join(visible), "".join(reasoning)
def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch):
# Parity contract: a reasoning-only reply must reach the client identically
# whether tools are on or off. Both generators stream <think> live then
# resolve to the bare reasoning text; the route's suffix-diff + extractor
# must therefore produce the same (visible, reasoning) split for both.
stream = [
_sse({"reasoning_content": "The capital"}),
_sse({"reasoning_content": " of France is Paris."}),
_done(),
]
tool_backend = _make_backend(monkeypatch, [list(stream)], [])
_patch_monotonic(monkeypatch, [1.0, 2.0, 2.0])
tool_cumulatives = [
e["text"]
for e in tool_backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "capital of France?"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
if e.get("type") == "content"
]
no_tool_backend = _make_backend(monkeypatch, [list(stream)], [])
no_tool_cumulatives = [
y
for y in no_tool_backend.generate_chat_completion(
messages = [{"role": "user", "content": "capital of France?"}],
)
if isinstance(y, str)
]
# Both paths stream the reasoning live with the same leading shape. (Raw
# yield lists aren't compared verbatim: the tool path emits a pre-existing
# duplicate trailing event that the route's suffix-diff dedupes.)
assert tool_cumulatives[:3] == no_tool_cumulatives[:3]
# The contract that matters: identical route-level output.
tool_out = _replay_route_reasoning_extractor(tool_cumulatives)
no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives)
assert tool_out == no_tool_out
# Pin the shared contract so a change to either path shows up here.
_visible, reasoning = tool_out
assert reasoning == "The capital of France is Paris."
def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch):
# _drain_silently sibling of the structured-tool close: a bare-JSON tool call
# with a live reasoning prefix must also close </think> before draining, and
# must never leak the drained call text as content.
tool_stream = [
_sse({"reasoning_content": "Searching now."}),
_sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}),
_done(),
]
final_stream = [
_sse({"content": "It is sunny."}),
_done(),
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads)
_patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0])
monkeypatch.setattr(
"core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny"
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "weather?"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start")
content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"]
assert content_before_tool[0] == "<think>Searching now."
assert content_before_tool[-1] == "<think>Searching now.</think>"
# The bare-JSON call text was drained, never surfaced as content.
assert not any('"name"' in t for t in content_before_tool)
def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch):
tool_stream = [

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

@ -1900,7 +1900,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
@ -2044,7 +2044,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
@ -2107,7 +2107,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
@ -2190,7 +2190,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
@ -2252,7 +2252,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
@ -2323,13 +2323,13 @@ class TestApiMonitorProviderAndCompletionStreams:
monitor_id = monitor_id,
)
)
await asyncio.wait_for(entered.wait(), timeout = 0.2)
await asyncio.wait_for(entered.wait(), timeout = 5.0)
assert cancel_id in inf_mod._CANCEL_REGISTRY
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
await asyncio.wait_for(cancelled.wait(), timeout = 0.2)
await asyncio.wait_for(cancelled.wait(), timeout = 5.0)
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
@ -2389,13 +2389,13 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
gate.set()
await asyncio.wait_for(returned.wait(), timeout = 0.2)
await asyncio.wait_for(returned.wait(), timeout = 5.0)
await asyncio.sleep(0)
await response._unstarted_cleanup()
assert upstream_response.is_closed

38
studio/backend/tests/test_safetensors_capability_advertise.py
View file

@ -48,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }}
"""
# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort
# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders
# identically to thinking-on-without-the-preamble), so the literal scan alone
# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to
# expose the encoder's full none/high/max ladder.
DEEPSEEK_V4_TEMPLATE = (
"{%- if not thinking is defined %}"
"{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}"
"{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n"
"{%- if thinking and reasoning_effort == 'max' %}"
"{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n"
"{%- for message in messages %}{{- message.content }}{%- endfor %}"
)
PLAIN_TEMPLATE = """
{%- for message in messages %}
{{- message.role + ': ' + message.content + '\\n' }}
@ -90,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false():
assert flags["reasoning_style"] == "enable_thinking"
def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max():
"""DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble.
Classified as the hybrid style with the full none/high/max ladder even
though the template only branches on 'max'."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF")
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
assert flags["reasoning_always_on"] is False
def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected():
"""The 'high' injection is scoped to deepseek-v4: a different model whose
template only branches on 'max' keeps ['max'] (no phantom 'high')."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF")
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["max"]
def test_detect_safetensors_features_passes_template_through_to_classifier():
"""Route wrapper forwards a real template to the inner classifier."""
from routes.inference import _detect_safetensors_features

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

@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base():
if runs_gate and not resolves_base:
offenders.append(f"{rel} runs a load gate but never resolves the LoRA base")
assert not offenders, "\n".join(offenders)
def test_rag_embedding_path_runs_the_malware_gate():
"""The RAG embedding model is set through /settings and later loaded by
SentenceTransformer, which deserializes pickles; both sites must run the malware gate
or a flagged repo loads unscanned (bypassing the normal model-load protections)."""
offenders = []
for rel in ("routes/settings.py", "core/rag/embeddings.py"):
if "evaluate_file_security(" not in (_BACKEND / rel).read_text():
offenders.append(
f"{rel} loads/persists an embedding model without evaluate_file_security"
)
assert not offenders, "\n".join(offenders)

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

@ -487,7 +487,7 @@ TEMPLATE_TO_RESPONSES_MAPPER = {
},
"qwen3-thinking": {
"instruction": "<|im_start|>user\n",
"response": "<|im_start|>assistant\n<think>\n",
"response": "<|im_start|>assistant\n<think>",
},
"qwen3": {
"instruction": "<|im_start|>user\n",

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

@ -83,11 +83,6 @@ import {
ZapIcon,
} from "@hugeicons/core-free-icons";
import { TestTubeOutlineIcon } from "@/lib/hugeicons-derived";
import {
exportConversationRawJsonl,
exportConversationCsv,
exportConversationShareGPT,
} from "@/features/chat/prompt-storage/prompt-storage-dialog";
import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage";
import {
Tooltip,
@ -169,6 +164,35 @@ function getTourId(pathname: string): string | null {
return null;
}
type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl";
const CHAT_EXPORT_OPTIONS: Array<{
label: string;
format: ConversationExportFormat;
}> = [
{ label: "Raw JSONL", format: "raw-jsonl" },
{ label: "CSV", format: "csv" },
{ label: "ShareGPT JSONL", format: "sharegpt-jsonl" },
];
async function exportConversationByFormat(
threadId: string,
format: ConversationExportFormat,
): Promise<void> {
const exports = await import(
"@/features/chat/prompt-storage/prompt-storage-dialog"
);
switch (format) {
case "raw-jsonl":
return exports.exportConversationRawJsonl(threadId);
case "csv":
return exports.exportConversationCsv(threadId);
case "sharegpt-jsonl":
return exports.exportConversationShareGPT(threadId);
}
}
function runStatusDotClass(status: TrainingRunSummary["status"]): string {
switch (status) {
case "running":
@ -354,7 +378,11 @@ export function AppSidebar() {
const activeProjectId = isChatRoute
? ((search.project as string | undefined) ?? null)
: null;
const { items: allChatItems } = useChatSidebarItems({
const {
items: allChatItems,
archivedItems: archivedChatItems,
loaded: chatItemsLoaded,
} = useChatSidebarItems({
enabled: !isStudioRoute,
requireMessages: false,
});
@ -890,11 +918,7 @@ export function AppSidebar() {
<span>Export</span>
</DropdownMenuSubTrigger>
<DropdownMenuSubContent sideOffset={8} alignOffset={-4} className="unsloth-plus-menu w-52">
{[
{ label: "Raw JSONL", fn: exportConversationRawJsonl },
{ label: "CSV", fn: exportConversationCsv },
{ label: "ShareGPT JSONL", fn: exportConversationShareGPT },
].map(({ label, fn }) => (
{CHAT_EXPORT_OPTIONS.map(({ label, format }) => (
<DropdownMenuItem
key={label}
onSelect={async () => {
@ -902,7 +926,9 @@ export function AppSidebar() {
const ids = item.type === "single"
? [item.id]
: (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id);
await Promise.all(ids.map((id) => fn(id)));
await Promise.all(
ids.map((id) => exportConversationByFormat(id, format)),
);
} catch {
toast.error("Export failed.");
}
@ -1324,6 +1350,16 @@ export function AppSidebar() {
renderChatSidebarItem(item, "recent"),
)}
</SidebarMenu>
{/* "No chats yet" only when there is truly no history:
project-scoped and archived threads leave Recents empty
but still count as existing chats. */}
{chatItemsLoaded &&
allChatItems.length === 0 &&
archivedChatItems.length === 0 && (
<p className="px-3 py-2 text-xs text-muted-foreground">
{t("shell.navigation.noChatsYet")}
</p>
)}
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>

483
studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx Normal file
View file

@ -0,0 +1,483 @@
// 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 {
Sheet,
SheetContent,
SheetDescription,
SheetHeader,
SheetTitle,
} from "@/components/ui/sheet";
import {
customProviderDisplayName,
parseExternalModelId,
useChatPreferencesStore,
useChatRuntimeStore,
useExternalProvidersStore,
} from "@/features/chat";
import { cn } from "@/lib/utils";
import { FileDatabaseIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useMessage, useMessageTiming } from "@assistant-ui/react";
import type { FC, ReactNode } from "react";
type ResponseDetailsMetadata = {
modelId?: string;
modelLabel?: string;
responseModelId?: string;
providerId?: string;
providerName?: string;
providerType?: string;
startedAt?: number;
finishedAt?: number;
durationMs?: number;
sessionId?: string | null;
cancelId?: string;
toolCalls?: string[];
tools?: Record<string, boolean | undefined>;
};
type ContextUsageMetadata = {
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
cachedTokens?: number;
cacheWriteTokens?: number;
modelId?: string;
};
type MessageCustomMetadata = {
responseDetails?: ResponseDetailsMetadata;
contextUsage?: ContextUsageMetadata;
serverTimings?: Record<string, unknown>;
reasoningDuration?: number;
};
function asNumber(value: unknown): number | undefined {
return typeof value === "number" && Number.isFinite(value)
? value
: undefined;
}
function formatNumber(value: number | undefined): string | null {
return value == null ? null : value.toLocaleString();
}
function formatMs(value: number | undefined): string | null {
if (value == null) return null;
if (value < 1000) return `${Math.round(value)}ms`;
return `${(value / 1000).toFixed(2)}s`;
}
function formatRate(value: number | undefined): string | null {
if (value == null) return null;
return `${value.toFixed(1)} tok/s`;
}
function formatDate(value: Date | number | string | undefined): string | null {
if (value == null) return null;
const date = value instanceof Date ? value : new Date(value);
if (Number.isNaN(date.getTime())) return null;
return new Intl.DateTimeFormat(undefined, {
dateStyle: "medium",
timeStyle: "medium",
}).format(date);
}
const TOOL_CATEGORY_LABELS: Record<string, string> = {
search: "Search",
fetch: "Fetch",
code: "Code",
images: "Images",
mcp: "MCP",
docs: "Docs",
artifacts: "Canvas",
};
const TOOL_CALL_LABELS: Record<string, string> = {
web_search: "Search",
web_fetch: "Fetch",
code_execution: "Code",
python: "Python",
terminal: "Terminal",
image_generation: "Images",
search_knowledge_base: "Docs",
render_html: "Canvas",
};
function uniqueValues(values: string[]): string[] {
return Array.from(new Set(values));
}
function toolCategoryFromCall(toolName: string): string | null {
const normalized = toolName.toLowerCase();
if (normalized === "web_search") return "search";
if (normalized === "web_fetch") return "fetch";
if (
normalized === "code_execution" ||
normalized === "python" ||
normalized === "terminal"
) {
return "code";
}
if (normalized === "image_generation") return "images";
if (normalized === "search_knowledge_base") return "docs";
if (normalized === "render_html") return "artifacts";
if (normalized.startsWith("mcp__")) return "mcp";
return null;
}
function formatToolCallName(toolName: string): string {
const normalized = toolName.toLowerCase();
if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized];
if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`;
return toolName
.replace(/[_-]+/g, " ")
.replace(/\b\w/g, (letter) => letter.toUpperCase());
}
function toolCallsFromContent(content: unknown): string[] {
if (!Array.isArray(content)) return [];
return uniqueValues(
content
.map((part) =>
part && typeof part === "object" && "type" in part
? (part as { type?: unknown; toolName?: unknown })
: null,
)
.filter(
(part): part is { type: "tool-call"; toolName: string } =>
part?.type === "tool-call" &&
typeof part.toolName === "string" &&
part.toolName.length > 0,
)
.map((part) => part.toolName),
);
}
function enabledTools(
tools: Record<string, boolean | undefined> | undefined,
toolCalls: string[],
): string | null {
if (!tools && toolCalls.length === 0) return null;
const activeKeys = new Set<string>();
for (const key of Object.keys(TOOL_CATEGORY_LABELS)) {
if (tools?.[key] === true) activeKeys.add(key);
}
for (const toolName of toolCalls) {
const key = toolCategoryFromCall(toolName);
if (key) activeKeys.add(key);
}
const active = Object.keys(TOOL_CATEGORY_LABELS)
.filter((key) => activeKeys.has(key))
.map((key) => TOOL_CATEGORY_LABELS[key]);
return active.length > 0 ? active.join(", ") : "None";
}
function calledTools(toolCalls: string[]): string | null {
if (toolCalls.length === 0) return null;
return uniqueValues(toolCalls.map(formatToolCallName)).join(", ");
}
function DetailSection({
title,
children,
}: {
title: string;
children: ReactNode;
}) {
return (
<section className="rounded-md bg-muted/45 p-3">
<h3 className="mb-2 font-heading text-foreground text-sm">{title}</h3>
<div className="grid gap-2">{children}</div>
</section>
);
}
function DetailRow({
label,
value,
mono = false,
}: {
label: string;
value: ReactNode | null | undefined;
mono?: boolean;
}) {
if (value == null || value === "") return null;
return (
<div className="grid grid-cols-[8.5rem_minmax(0,1fr)] items-start gap-3 text-[13px]">
<span className="text-muted-foreground">{label}</span>
<span
className={cn(
"min-w-0 break-words text-right text-foreground",
mono && "font-mono tabular-nums",
)}
>
{value}
</span>
</div>
);
}
function useResponseModelDisplay() {
const message = useMessage();
const models = useChatRuntimeStore((s) => s.models);
const providers = useExternalProvidersStore((s) => s.providers);
const custom = (
message.metadata as Record<string, unknown> | undefined
)?.custom as MessageCustomMetadata | undefined;
const responseDetails = custom?.responseDetails;
const usage = custom?.contextUsage;
const serverTimings = custom?.serverTimings;
const recordedModelId =
responseDetails?.responseModelId ??
responseDetails?.modelId ??
usage?.modelId;
const parsedExternal = parseExternalModelId(recordedModelId);
const provider = parsedExternal
? providers.find((candidate) => candidate.id === parsedExternal.providerId)
: null;
const modelSummary = models.find(
(candidate) => candidate.id === recordedModelId,
);
const modelLabel =
responseDetails?.modelLabel ??
responseDetails?.responseModelId ??
parsedExternal?.modelId ??
modelSummary?.name ??
recordedModelId ??
"Not recorded";
const providerLabel =
responseDetails?.providerName ??
provider?.name ??
(responseDetails?.providerType
? customProviderDisplayName(responseDetails.providerType)
: parsedExternal
? customProviderDisplayName(provider?.providerType)
: recordedModelId
? "Local model"
: null);
return {
message,
custom,
responseDetails,
usage,
serverTimings,
modelLabel,
providerLabel,
};
}
export const MessageResponseModelBadge: FC<{ className?: string }> = ({
className,
}) => {
const showResponseModel = useChatPreferencesStore(
(state) => state.showResponseModel,
);
const { modelLabel, providerLabel } = useResponseModelDisplay();
if (!showResponseModel || modelLabel === "Not recorded") {
return null;
}
return (
<span
className={cn(
"aui-response-model-badge inline-flex min-h-5 max-w-full items-center text-muted-foreground/80 text-xs font-medium leading-5 opacity-0 transition-opacity duration-150 group-hover/assistant-message:opacity-100 group-focus-within/assistant-message:opacity-100",
className,
)}
title={providerLabel ? `${modelLabel} - ${providerLabel}` : modelLabel}
>
<span className="min-w-0 truncate align-middle">{modelLabel}</span>
</span>
);
};
export const MessageResponseDetailsSheet: FC<{
open: boolean;
onOpenChange: (open: boolean) => void;
}> = ({ open, onOpenChange }) => {
const timing = useMessageTiming();
const {
message,
responseDetails,
usage,
serverTimings,
modelLabel,
providerLabel,
} = useResponseModelDisplay();
const promptTokens =
usage?.promptTokens ?? asNumber(serverTimings?.prompt_n);
const completionTokens =
usage?.completionTokens ??
timing?.tokenCount ??
asNumber(serverTimings?.predicted_n);
const totalTokens =
usage?.totalTokens ??
(promptTokens != null && completionTokens != null
? promptTokens + completionTokens
: undefined);
const totalTime =
responseDetails?.durationMs ?? timing?.totalStreamTime ?? undefined;
const summaryLabel =
modelLabel === "Not recorded" ? "Model not recorded" : `Used ${modelLabel}`;
const messageToolCalls = toolCallsFromContent(message.content);
const toolCalls =
responseDetails?.toolCalls && responseDetails.toolCalls.length > 0
? responseDetails.toolCalls
: messageToolCalls;
return (
<Sheet open={open} onOpenChange={onOpenChange}>
<SheetContent
side="right"
className="w-[min(28rem,100vw)] p-0 sm:max-w-[28rem]"
>
<SheetHeader className="border-b p-4">
<SheetTitle className="flex items-center gap-2 pr-10 font-heading text-base">
<HugeiconsIcon
icon={FileDatabaseIcon}
strokeWidth={1.75}
className="size-icon text-chat-icon-fg"
/>
Response details
</SheetTitle>
<SheetDescription className="sr-only">
Timing, model, token, and tool details for this response.
</SheetDescription>
</SheetHeader>
<div className="flex min-h-0 flex-1 flex-col gap-4 overflow-y-auto p-4">
<div className="min-w-0 rounded-md border border-border/70 bg-card p-3">
<p className="min-w-0 break-words font-heading text-foreground text-sm">
{summaryLabel}
</p>
{providerLabel ? (
<p className="mt-1 min-w-0 break-words text-muted-foreground text-xs">
{providerLabel}
</p>
) : null}
</div>
<DetailSection title="Response">
<DetailRow label="Model" value={modelLabel} />
<DetailRow
label="Requested"
value={
responseDetails?.modelId &&
responseDetails.modelId !== responseDetails.responseModelId
? responseDetails.modelId
: null
}
/>
<DetailRow label="Provider" value={providerLabel} />
<DetailRow label="Message ID" value={message.id} mono={true} />
<DetailRow label="Created" value={formatDate(message.createdAt)} />
<DetailRow
label="Started"
value={formatDate(responseDetails?.startedAt)}
/>
<DetailRow
label="Finished"
value={formatDate(responseDetails?.finishedAt)}
/>
</DetailSection>
<DetailSection title="Tokens">
<DetailRow label="Prompt" value={formatNumber(promptTokens)} mono />
<DetailRow
label="Output"
value={formatNumber(completionTokens)}
mono
/>
<DetailRow label="Total" value={formatNumber(totalTokens)} mono />
<DetailRow
label="Cache hits"
value={formatNumber(
usage?.cachedTokens ?? asNumber(serverTimings?.cache_n),
)}
mono
/>
<DetailRow
label="Cache writes"
value={formatNumber(usage?.cacheWriteTokens)}
mono
/>
</DetailSection>
<DetailSection title="Timing">
<DetailRow label="Total" value={formatMs(totalTime)} mono />
<DetailRow
label="First token"
value={formatMs(timing?.firstTokenTime)}
mono
/>
<DetailRow
label="Prompt eval"
value={formatMs(asNumber(serverTimings?.prompt_ms))}
mono
/>
<DetailRow
label="Generation"
value={formatMs(asNumber(serverTimings?.predicted_ms))}
mono
/>
<DetailRow
label="Speed"
value={formatRate(
asNumber(serverTimings?.predicted_per_second) ??
timing?.tokensPerSecond,
)}
mono
/>
<DetailRow
label="Chunks"
value={formatNumber(timing?.totalChunks)}
mono
/>
<DetailRow
label="Tool calls"
value={formatNumber(timing?.toolCallCount)}
mono
/>
</DetailSection>
<DetailSection title="Tools">
<DetailRow
label="Enabled"
value={enabledTools(responseDetails?.tools, toolCalls)}
/>
<DetailRow label="Called" value={calledTools(toolCalls)} />
<DetailRow
label="Confirmation"
value={
responseDetails?.tools?.confirmToolCalls === true
? "On"
: responseDetails?.tools?.confirmToolCalls === false
? "Off"
: null
}
/>
<DetailRow
label="Bypass"
value={
responseDetails?.tools?.bypassPermissions === true
? "On"
: responseDetails?.tools?.bypassPermissions === false
? "Off"
: null
}
/>
<DetailRow label="Session" value={responseDetails?.sessionId} mono />
<DetailRow label="Run ID" value={responseDetails?.cancelId} mono />
</DetailSection>
</div>
</SheetContent>
</Sheet>
);
};

10
studio/frontend/src/components/assistant-ui/reasoning.tsx
View file

@ -6,6 +6,7 @@
/* eslint-disable react-refresh/only-export-components */
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { MessageResponseModelBadge } from "@/components/assistant-ui/message-response-details-sheet";
import {
Collapsible,
CollapsibleContent,
@ -390,14 +391,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
onOpenChange={handleOpenChange}
variant={variant}
>
<div className="flex items-center gap-2">
<div className="flex min-w-0 items-center gap-2">
<ReasoningTrigger
className="min-w-0 flex-1"
className="min-w-0 flex-none"
active={isReasoningStreaming}
// Prefer server timing when available.
duration={persistedDuration || duration}
/>
<div className="flex w-16 shrink-0 justify-end">
<span className="hidden min-w-0 max-w-[12rem] group-hover/assistant-message:inline-flex group-focus-within/assistant-message:inline-flex sm:max-w-[16rem]">
<MessageResponseModelBadge className="min-w-0" />
</span>
<div className="ml-auto flex w-16 shrink-0 justify-end">
{isOpen && !isReasoningStreaming && (
<ReasoningCopyButton startIndex={startIndex} endIndex={endIndex} />
)}

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

@ -12,6 +12,10 @@ import {
import { downloadImagePart } from "@/components/assistant-ui/image";
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts";
import {
MessageResponseDetailsSheet,
MessageResponseModelBadge,
} from "@/components/assistant-ui/message-response-details-sheet";
import { MessageTiming } from "@/components/assistant-ui/message-timing";
import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources";
@ -3564,6 +3568,9 @@ const AssistantMessage: FC = () => {
const aui = useAui();
const messageId = useAuiState(({ message }) => message.id);
const messageContent = useAuiState(({ message }) => message.content);
const hasReasoningParts = useAuiState(({ message }) =>
message.parts.some((part) => part.type === "reasoning"),
);
const incognito = useChatRuntimeStore((s) => s.incognito);
// Use global store for editing state to ensure a single source of truth
@ -3620,7 +3627,7 @@ const AssistantMessage: FC = () => {
return (
<MessagePrimitive.Root
className="aui-assistant-message-root relative mx-auto min-w-0 w-full max-w-(--thread-content-max-width) pt-0.5 pb-4 text-[15.5px] [font-weight:410] tracking-[0.01em] dark:tracking-[0.02em]"
className="group/assistant-message aui-assistant-message-root relative mx-auto min-w-0 w-full max-w-(--thread-content-max-width) pt-0.5 pb-4 text-[15.5px] [font-weight:410] tracking-[0.01em] dark:tracking-[0.02em]"
data-role="assistant"
>
<div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
@ -3649,6 +3656,11 @@ const AssistantMessage: FC = () => {
</div>
) : (
<>
{!hasReasoningParts ? (
<div className="pointer-events-none relative h-0 min-w-0">
<MessageResponseModelBadge className="absolute -top-6 left-0 max-w-[min(22rem,100%)]" />
</div>
) : null}
<GeneratingIndicator />
<CancelledIndicator />
<DiffusionCanvas />
@ -3893,58 +3905,76 @@ const EditAssistantMessageButton: FC = () => {
const AssistantActionBar: FC = () => {
const { forkMessage, forkDisabled } = useForkMessageAction();
const [detailsOpen, setDetailsOpen] = useState(false);
return (
<ActionBarPrimitive.Root
hideWhenRunning={true}
className="aui-assistant-action-bar-root col-start-3 row-start-2 flex items-center gap-1 text-chat-icon-fg [&_button:not([data-slot=message-timing-trigger])]:size-8 [&_button]:!rounded-full [&_button:hover]:bg-chat-icon-bg-hover [&_button:hover]:text-chat-icon-fg-hover"
>
<CopyButton />
<EditAssistantMessageButton />
<ActionBarPrimitive.Reload asChild={true}>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon strokeWidth={1.75} className="size-icon" />
</TooltipIconButton>
</ActionBarPrimitive.Reload>
<ForkCountBadge />
<DeleteMessageButton />
<ActionBarMorePrimitive.Root>
<ActionBarMorePrimitive.Trigger asChild={true}>
<TooltipIconButton
tooltip="More"
className="data-[state=open]:bg-accent"
>
<MoreHorizontalIcon strokeWidth={1.75} className="size-icon" />
<>
<ActionBarPrimitive.Root
hideWhenRunning={true}
className="aui-assistant-action-bar-root col-start-3 row-start-2 flex items-center gap-1 text-chat-icon-fg [&_button:not([data-slot=message-timing-trigger])]:size-8 [&_button]:!rounded-full [&_button:hover]:bg-chat-icon-bg-hover [&_button:hover]:text-chat-icon-fg-hover"
>
<CopyButton />
<EditAssistantMessageButton />
<ActionBarPrimitive.Reload asChild={true}>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon strokeWidth={1.75} className="size-icon" />
</TooltipIconButton>
</ActionBarMorePrimitive.Trigger>
<ActionBarMorePrimitive.Content
side="bottom"
align="start"
onCloseAutoFocus={(e) => e.preventDefault()}
className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
>
<ActionBarMorePrimitive.Item
disabled={forkDisabled}
onSelect={() => void forkMessage()}
className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
</ActionBarPrimitive.Reload>
<ForkCountBadge />
<DeleteMessageButton />
<ActionBarMorePrimitive.Root>
<ActionBarMorePrimitive.Trigger asChild={true}>
<TooltipIconButton
tooltip="More"
className="data-[state=open]:bg-accent"
>
<MoreHorizontalIcon strokeWidth={1.75} className="size-icon" />
</TooltipIconButton>
</ActionBarMorePrimitive.Trigger>
<ActionBarMorePrimitive.Content
side="bottom"
align="start"
onCloseAutoFocus={(e) => e.preventDefault()}
className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
>
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
Fork in new chat
</ActionBarMorePrimitive.Item>
<ActionBarPrimitive.ExportMarkdown asChild={true}>
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
<ActionBarMorePrimitive.Item
disabled={forkDisabled}
onSelect={() => void forkMessage()}
className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
>
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
Fork in new chat
</ActionBarMorePrimitive.Item>
<ActionBarPrimitive.ExportMarkdown asChild={true}>
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 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>
</ActionBarPrimitive.ExportMarkdown>
<ActionBarMorePrimitive.Item
onSelect={() => setDetailsOpen(true)}
className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground"
>
<HugeiconsIcon
icon={Download01Icon}
icon={FileDatabaseIcon}
strokeWidth={1.75}
className="size-icon"
/>
Export as Markdown
See response details
</ActionBarMorePrimitive.Item>
</ActionBarPrimitive.ExportMarkdown>
</ActionBarMorePrimitive.Content>
</ActionBarMorePrimitive.Root>
<MessageTiming side="top" className="h-8 px-2" />
</ActionBarPrimitive.Root>
</ActionBarMorePrimitive.Content>
</ActionBarMorePrimitive.Root>
<MessageTiming side="top" className="h-8 px-2" />
</ActionBarPrimitive.Root>
<MessageResponseDetailsSheet
open={detailsOpen}
onOpenChange={setDetailsOpen}
/>
</>
);
};

26
studio/frontend/src/config/env.ts
View file

@ -54,6 +54,17 @@ export const usePlatformStore = create<PlatformState>()((_, get) => ({
isChatOnly: () => get().chatOnly,
}));
// Once an authoritative (server-reported) platform has been fetched, a
// non-forced response must not overwrite it. The post-render fetchDeviceType()
// in main.tsx runs before auth is ready and can resolve after the authed
// root-route/provider fetches; such a late write would reset deviceType,
// cloudflareUrl/serverUrl/secure, and fetched, whether it is a browser fallback
// (unauthenticated) or an earlier authenticated request that landed after a
// later forced refresh. Forced refreshes are explicit re-reads, so they still write.
function shouldKeepAuthoritativePlatform(force?: boolean): boolean {
return !force && usePlatformStore.getState().fetched;
}
// `force` re-reads /api/health even if cached, to pick up a late-arriving tunnel URL.
export async function fetchDeviceType(options?: {
force?: boolean;
@ -81,6 +92,15 @@ export async function fetchDeviceType(options?: {
server_url?: string | null;
secure?: boolean;
};
// Once the store holds an authoritative (server-reported) platform, a
// non-forced response must not overwrite it. It may be an unauthenticated
// fallback, or an earlier authenticated request that resolved after a
// later forced refresh already picked up device_type and the tunnel
// fields; writing either would reset device type or null the tunnel
// fields. Forced refreshes are explicit re-reads, so they still write.
if (shouldKeepAuthoritativePlatform(options?.force)) {
return usePlatformStore.getState().deviceType;
}
const deviceType = data.device_type ?? detectLocalPlatform();
const chatOnly = data.chat_only ?? false;
const chatOnlyReason = data.chat_only_reason ?? null;
@ -101,7 +121,11 @@ export async function fetchDeviceType(options?: {
} catch {
// Backend not ready: use client-side detection so chat-only guard works
// on initial load (important for macOS). Keep fetched=false so a later
// call retries against the backend.
// call retries against the backend. But a late non-forced failure must not
// wipe an authoritative platform that already resolved.
if (shouldKeepAuthoritativePlatform(options?.force)) {
return usePlatformStore.getState().deviceType;
}
const deviceType = detectLocalPlatform();
const chatOnly = deviceType === "mac";
usePlatformStore.setState({ deviceType, chatOnly, fetched: false });

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

@ -140,6 +140,32 @@ interface ServerTimings {
diffusion_steps_per_second?: number;
}
interface ResponseDetailsMetadata {
modelId: string;
modelLabel: string;
responseModelId: string;
providerId?: string;
providerName: string;
providerType: string;
startedAt: number;
finishedAt: number;
durationMs: number;
sessionId?: string;
cancelId: string;
toolCalls: string[];
tools: {
search: boolean;
fetch: boolean;
code: boolean;
images: boolean;
mcp: boolean;
docs: boolean;
artifacts: boolean;
confirmToolCalls: boolean;
bypassPermissions: boolean;
};
}
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
type RunMessage = RunMessages[number];
@ -1769,6 +1795,9 @@ export function createOpenAIStreamAdapter(
(provider) => provider.id === externalSelection.providerId,
)
: null;
const selectedModelSummary = runtime.models.find(
(model) => model.id === params.checkpoint,
);
const externalApiKey = externalProvider
? getExternalProviderApiKey(externalProvider.id).trim()
: "";
@ -2151,6 +2180,7 @@ export function createOpenAIStreamAdapter(
let waitingFirstChunk = true;
let firstTokenSettled = false;
const streamStartTime = Date.now();
let responseModelId = externalSelection?.modelId ?? params.checkpoint;
let firstTokenTime: number | undefined;
let totalChunks = 0;
let resolveFirstToken: (() => void) | null = null;
@ -2372,6 +2402,59 @@ export function createOpenAIStreamAdapter(
const externalBackendProviderType = toExternalBackendProviderType(
externalProvider?.providerType,
);
const buildResponseDetails = (
finishedAt: number,
): ResponseDetailsMetadata => ({
modelId: params.checkpoint,
modelLabel:
(isExternalRequest || responseModelId !== params.checkpoint
? responseModelId
: selectedModelSummary?.name || responseModelId) ||
params.checkpoint ||
"Unknown model",
responseModelId:
responseModelId ||
externalSelection?.modelId ||
params.checkpoint,
...(externalProvider?.id ? { providerId: externalProvider.id } : {}),
providerName:
externalProvider?.name ??
(isExternalRequest ? "External provider" : "Local model"),
providerType: externalProvider?.providerType ?? "local",
startedAt: streamStartTime,
finishedAt,
durationMs: finishedAt - streamStartTime,
...(sandboxSessionId ? { sessionId: sandboxSessionId } : {}),
cancelId,
toolCalls: Array.from(
new Set(
toolCallParts
.map((part) => part.toolName)
.filter(
(toolName): toolName is string =>
typeof toolName === "string" && toolName.length > 0,
),
),
),
tools: {
search:
webSearchEnabledForThisTurn ||
(!isExternalRequest && supportsTools && toolsEnabled),
fetch: webFetchEnabledForThisTurn,
code:
codeExecEnabledForThisTurn ||
(!isExternalRequest && supportsTools && codeToolsEnabled),
images: imageGenerationEnabledForThisTurn,
mcp: !isExternalRequest && supportsTools && mcpEnabledForChat,
docs:
!isExternalRequest &&
supportsTools &&
(ragEnabled || projectRagEnabled),
artifacts: renderHtmlToolEnabledForThisTurn,
confirmToolCalls,
bypassPermissions,
},
});
const externalCapabilities = getProviderCapabilities(
externalProvider?.providerType,
);
@ -2768,6 +2851,11 @@ export function createOpenAIStreamAdapter(
const stream = streamChatCompletions(requestPayload, abortSignal);
for await (const chunk of stream) {
const chunkModel = (chunk as { model?: unknown }).model;
if (typeof chunkModel === "string" && chunkModel.length > 0) {
responseModelId = chunkModel;
}
// Handle tool status events
const toolStatusText = (
chunk as unknown as { _toolStatus?: string }
@ -3435,11 +3523,12 @@ export function createOpenAIStreamAdapter(
});
}
const finishedAt = Date.now();
const finalTiming = buildTiming(
streamStartTime,
totalChunks,
serverPromptEvalTime ?? firstTokenTime,
Date.now() - streamStartTime,
finishedAt - streamStartTime,
finalTokenCount,
toolCallParts.length,
finalTokPerSec,
@ -3475,6 +3564,7 @@ export function createOpenAIStreamAdapter(
modelId: params.checkpoint,
}
: undefined,
responseDetails: buildResponseDetails(finishedAt),
timing: finalTiming,
},
},

43
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -35,7 +35,6 @@ import {
useNativeModelDrop,
useNativePathLeasesSupported,
} from "@/features/native-intents";
import { ProjectSourcesPanel } from "@/features/rag/components/project-sources-panel";
import { GuidedTour, useGuidedTourController } from "@/features/tour";
import { isTauri } from "@/lib/api-base";
import { toast } from "@/lib/toast";
@ -51,7 +50,9 @@ import { Tooltip as TooltipPrimitive } from "radix-ui";
import {
type CSSProperties,
type ReactElement,
lazy,
memo,
Suspense,
useCallback,
useEffect,
useMemo,
@ -134,6 +135,13 @@ import {
} from "./utils/chat-history-storage";
import { isAssistantLocalThreadId } from "./utils/thread-ids";
const ProjectSourcesPanel = lazy(() =>
import("@/features/rag/components/project-sources-panel").then((module) => ({
default: module.ProjectSourcesPanel,
})),
);
type LoraCandidate = {
id: string;
baseModel: string;
@ -1018,7 +1026,15 @@ function ProjectLanding({
</div>
{projectTab === "sources" ? (
<ProjectSourcesPanel projectId={projectId} />
<Suspense
fallback={
<div className="mt-8 rounded-[26px] bg-muted/30 px-6 py-10 text-center text-sm text-muted-foreground">
Loading sources…
</div>
}
>
<ProjectSourcesPanel projectId={projectId} />
</Suspense>
) : (
<div className="mt-8 flex flex-col gap-1">
{items.map((item) => {
@ -2246,12 +2262,29 @@ export function ChatPage({
return [...fromLoras, ...localModels];
}, [lorasFromStore, localModels]);
useEffect(() => {
if (getTrainingCompareHandoff()) return;
void refresh();
const inventoryRefreshStartedRef = useRef(false);
const refreshDeferredModelInventories = useCallback(() => {
inventoryRefreshStartedRef.current = true;
void refresh({ includeLoras: true });
refreshLocalModels();
}, [refresh, refreshLocalModels]);
useEffect(() => {
if (getTrainingCompareHandoff()) return;
void refresh({ includeLoras: false });
const timeoutId = window.setTimeout(() => {
if (!inventoryRefreshStartedRef.current) {
refreshDeferredModelInventories();
}
}, 1200);
return () => window.clearTimeout(timeoutId);
}, [refresh, refreshDeferredModelInventories]);
useEffect(() => {
if (!active || !modelSelectorOpen) return;
refreshDeferredModelInventories();
}, [active, modelSelectorOpen, refreshDeferredModelInventories]);
useEffect(() => {
// ChatPage no longer remounts on navigation, so re-check the handoff whenever
// we return to /chat (e.g. from the training progress "compare in chat" action).

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

@ -312,14 +312,18 @@ export function useChatModelRuntime() {
[],
);
const refresh = useCallback(async (options?: { signal?: AbortSignal }) => {
const refresh = useCallback(async (options?: {
signal?: AbortSignal;
includeLoras?: boolean;
}) => {
const signal = options?.signal;
const includeLoras = options?.includeLoras ?? true;
setModelsError(null);
try {
const [listRes, statusRes, lorasRes] = await Promise.all([
listModels(),
getInferenceStatus(),
listLoras(),
includeLoras ? listLoras() : Promise.resolve(null),
]);
// Cancellation can land while the requests above are in flight. Bail
@ -327,7 +331,9 @@ export function useChatModelRuntime() {
if (signal?.aborted) return;
setModels(listRes.models.map(toChatModelSummary));
setLoras(lorasRes.loras.map(toLoraSummary));
if (lorasRes) {
setLoras(lorasRes.loras.map(toLoraSummary));
}
const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint;
const isExternalSelectionActive = isExternalModelId(selectedCheckpoint);

26
studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
View file

@ -27,16 +27,21 @@ export interface SidebarItem {
id: string;
title: string;
createdAt: number;
updatedAt: number;
isFork?: boolean;
projectId?: string | null;
}
function lastActivityAt(thread: ThreadRecord): number {
return thread.updatedAt ?? thread.createdAt;
}
export function groupThreads(
threads: ThreadRecord[],
archived = false,
): SidebarItem[] {
const items: SidebarItem[] = [];
const seenPairs = new Set<string>();
const pairItems = new Map<string, SidebarItem>();
for (const t of threads) {
// Coerce archived to a boolean before comparing. Legacy threads (from the
@ -48,30 +53,35 @@ export function groupThreads(
continue;
}
if (t.pairId) {
if (seenPairs.has(t.pairId)) {
const existing = pairItems.get(t.pairId);
if (existing) {
existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t));
continue;
}
seenPairs.add(t.pairId);
items.push({
const item: SidebarItem = {
type: "compare",
id: t.pairId,
title: t.title,
createdAt: t.createdAt,
updatedAt: lastActivityAt(t),
projectId: t.projectId ?? null,
});
};
pairItems.set(t.pairId, item);
items.push(item);
} else if (!t.pairId) {
items.push({
type: "single",
id: t.id,
title: t.title,
createdAt: t.createdAt,
updatedAt: lastActivityAt(t),
isFork: Boolean(t.forkedFromThreadId),
projectId: t.projectId ?? null,
});
}
}
return items.sort((a, b) => b.createdAt - a.createdAt);
return items.sort((a, b) => b.updatedAt - a.updatedAt);
}
// Streaming fires CHAT_HISTORY_UPDATED_EVENT per chunk. Debounce so each quiet
@ -84,6 +94,7 @@ export function useChatSidebarItems(options?: {
requireMessages?: boolean;
}) {
const [allThreads, setAllThreads] = useState<ThreadRecord[]>([]);
const [loaded, setLoaded] = useState(false);
const enabled = options?.enabled ?? true;
const requireMessages = options?.requireMessages ?? true;
@ -111,6 +122,7 @@ export function useChatSidebarItems(options?: {
// were in flight, or if the effect was torn down.
if (cancelled || seq !== requestSeq) return;
setAllThreads(threads);
setLoaded(true);
} catch (error) {
if (isExpectedBackgroundChatStorageError(error)) {
return;
@ -144,7 +156,7 @@ export function useChatSidebarItems(options?: {
const archivedItems = groupThreads(allThreads ?? [], true);
const canCompare = useChatRuntimeStore((s) => Boolean(s.params.checkpoint));
return { items, archivedItems, canCompare };
return { items, archivedItems, canCompare, loaded };
}
function cancelIfRunning(threadId: string): void {

7
studio/frontend/src/features/chat/index.ts
View file

@ -26,7 +26,12 @@ export {
type PlusMenuItemId,
} from "./stores/plus-menu-prefs-store";
export { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
export { isExternalModelId } from "./external-providers";
export {
customProviderDisplayName,
isExternalModelId,
parseExternalModelId,
} from "./external-providers";
export { useExternalProvidersStore } from "./stores/external-providers-store";
export { ChatSearchDialog } from "./components/chat-search-dialog";
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
export type { ProjectRecord } from "./types";

12
studio/frontend/src/features/chat/runtime-provider.tsx
View file

@ -22,7 +22,6 @@ import {
unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime,
} from "@assistant-ui/react";
import { createAssistantStream } from "assistant-stream";
import mammoth from "mammoth";
import {
type ReactElement,
type ReactNode,
@ -33,7 +32,6 @@ import {
useMemo,
useRef,
} from "react";
import { extractText, getDocumentProxy } from "unpdf";
import { toast } from "sonner";
import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter";
import {
@ -181,7 +179,10 @@ class PDFAttachmentAdapter implements AttachmentAdapter {
}
async send(attachment: PendingAttachment): Promise<CompleteAttachment> {
const buffer = new Uint8Array(await attachment.file.arrayBuffer());
const [{ extractText, getDocumentProxy }, buffer] = await Promise.all([
import("unpdf"),
attachment.file.arrayBuffer().then((bytes) => new Uint8Array(bytes)),
]);
const pdf = await getDocumentProxy(buffer);
const { text } = await extractText(pdf, { mergePages: true });
return {
@ -298,7 +299,10 @@ class DocxAttachmentAdapter implements AttachmentAdapter {
}
async send(attachment: PendingAttachment): Promise<CompleteAttachment> {
const arrayBuffer = await attachment.file.arrayBuffer();
const [{ default: mammoth }, arrayBuffer] = await Promise.all([
import("mammoth"),
attachment.file.arrayBuffer(),
]);
const { value } = await mammoth.extractRawText({ arrayBuffer });
return {
id: attachment.id,

7
studio/frontend/src/features/chat/stores/chat-preferences-store.ts
View file

@ -7,11 +7,14 @@ import { persist } from "zustand/middleware";
// Client-side chat UI prefs kept in localStorage, not the chat DB.
// confirmDeleteChats: when off, deleting a chat skips the confirm dialog.
// showModelDisclaimer: when off, hide the "LLMs can make mistakes" footer note.
// showResponseModel: when on, assistant responses show the producing model.
export interface ChatPreferencesState {
confirmDeleteChats: boolean;
setConfirmDeleteChats: (value: boolean) => void;
showModelDisclaimer: boolean;
setShowModelDisclaimer: (value: boolean) => void;
showResponseModel: boolean;
setShowResponseModel: (value: boolean) => void;
}
export const useChatPreferencesStore = create<ChatPreferencesState>()(
@ -23,6 +26,9 @@ export const useChatPreferencesStore = create<ChatPreferencesState>()(
showModelDisclaimer: true,
setShowModelDisclaimer: (showModelDisclaimer) =>
set({ showModelDisclaimer }),
showResponseModel: false,
setShowResponseModel: (showResponseModel) =>
set({ showResponseModel }),
}),
{
name: "unsloth_chat_preferences",
@ -32,6 +38,7 @@ export const useChatPreferencesStore = create<ChatPreferencesState>()(
...current,
confirmDeleteChats: saved?.confirmDeleteChats ?? true,
showModelDisclaimer: saved?.showModelDisclaimer ?? true,
showResponseModel: saved?.showResponseModel ?? false,
};
},
},

1
studio/frontend/src/features/chat/types.ts
View file

@ -36,6 +36,7 @@ export interface ThreadRecord {
projectId?: string | null;
archived: boolean;
createdAt: number;
updatedAt?: number;
/**
* OpenAI shell tool container id from a prior response. When set, the
* next turn reuses it via `environment.type="container_reference"` so

5
studio/frontend/src/features/chat/utils/chat-history-storage.ts
View file

@ -590,7 +590,10 @@ export async function listStoredChatThreads(
}
return Array.from(byId.values())
.filter((thread) => matchesThreadListArgs(thread, args))
.sort((a, b) => b.createdAt - a.createdAt);
.sort(
(a, b) =>
(b.updatedAt ?? b.createdAt) - (a.updatedAt ?? a.createdAt),
);
}
export async function listStoredChatThreadsWithMessages(

9
studio/frontend/src/features/settings/api/embedding-model.ts
View file

@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = {
* (wrong type, gated repo, or offline). Retry with force to save anyway. */
export class EmbeddingModelVerificationError extends Error {}
/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan.
* A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */
export class EmbeddingModelBlockedError extends Error {}
function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings {
return {
embeddingModel: settings.embedding_model,
@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings(
force: options?.force ?? false,
}),
});
if (res.status === 403) {
throw new EmbeddingModelBlockedError(
await readFastApiError(res, "This model is blocked by a security scan"),
);
}
if (res.status === 409) {
throw new EmbeddingModelVerificationError(
await readFastApiError(res, "Could not verify the embedding model"),

15
studio/frontend/src/features/settings/tabs/chat-tab.tsx
View file

@ -213,6 +213,12 @@ export function ChatTab() {
const setShowModelDisclaimer = useChatPreferencesStore(
(state) => state.setShowModelDisclaimer,
);
const showResponseModel = useChatPreferencesStore(
(state) => state.showResponseModel,
);
const setShowResponseModel = useChatPreferencesStore(
(state) => state.setShowResponseModel,
);
useEffect(() => {
void countAllChats().then(setCount);
@ -412,6 +418,15 @@ export function ChatTab() {
onCheckedChange={setShowModelDisclaimer}
/>
</SettingsRow>
<SettingsRow
label="Show response model"
description="Show model metadata in assistant responses."
>
<Switch
checked={showResponseModel}
onCheckedChange={setShowResponseModel}
/>
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.chat.artifacts.title")}>

6
studio/frontend/src/features/settings/tabs/general-tab.tsx
View file

@ -42,6 +42,7 @@ import {
updatePreviewSharing,
} from "../api/preview-sharing";
import {
EmbeddingModelBlockedError,
type EmbeddingModelSettings,
EmbeddingModelVerificationError,
loadEmbeddingModelSettings,
@ -410,7 +411,10 @@ export function GeneralTab() {
description: t("settings.general.rag.reindexWarning"),
});
} catch (error) {
if (error instanceof EmbeddingModelVerificationError) {
// A hard security block cannot be forced; keep the "save anyway" action hidden.
if (error instanceof EmbeddingModelBlockedError) {
setEmbeddingModelNeedsForce(false);
} else if (error instanceof EmbeddingModelVerificationError) {
setEmbeddingModelNeedsForce(true);
}
setEmbeddingModelError(

1
studio/frontend/src/i18n/locales/en.ts
View file

@ -43,6 +43,7 @@ export const en = {
video: "Video",
export: "Export",
recents: "Recents",
noChatsYet: "No chats yet",
settings: "Settings",
api: "API",
lightMode: "Light Mode",

1
studio/frontend/src/i18n/locales/zh-CN.ts
View file

@ -43,6 +43,7 @@ export const zhCN = {
video: "视频",
export: "导出",
recents: "最近",
noChatsYet: "暂无对话",
settings: "设置",
api: "API",
lightMode: "浅色模式",

169
studio/frontend/src/lib/latex.ts
View file

@ -1,8 +1,13 @@
// Adapted from LibreChat's latex.ts
// https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts
//
// Escapes currency dollar signs so they are not misinterpreted as LaTeX math
// delimiters when singleDollarTextMath is enabled.
// Two jobs, in order:
// 1. Convert LaTeX bracket delimiters (`\[...\]`, `\(...\)`) into the dollar
// forms remark-math understands (`$$...$$`, `$...$`). remark-math only
// tokenizes dollar delimiters, so models that emit `\[...\]` / `\(...\)`
// would otherwise render as literal text.
// 2. Escape currency dollar signs so they are not misinterpreted as LaTeX
// math delimiters when singleDollarTextMath is enabled.
/**
* Matches a single $ followed by a number pattern (currency), e.g.:
@ -15,14 +20,14 @@ const CURRENCY_REGEX =
/(?<![\\$])\$(?!\$)(?=\d+(?:,\d{3})*(?:\.\d+)?[KMBkmb]?(?:\s|$|[^a-zA-Z\d]))/g;
/**
* Find code-block regions (``` ... ``` and ` ... `) to skip.
* Find code-block regions (``` ... ```, ~~~ ... ~~~, and ` ... `) to skip.
* Returns a sorted array of [start, end] index pairs.
*/
function findCodeBlockRegions(content: string): Array<[number, number]> {
const regions: Array<[number, number]> = [];
// Fenced code blocks: ```...```
const fencedRe = /```[\s\S]*?```/g;
// Fenced code blocks: ```...``` and ~~~...~~~ (both are code in GFM)
const fencedRe = /```[\s\S]*?```|~~~[\s\S]*?~~~/g;
let match: RegExpExecArray | null;
while ((match = fencedRe.exec(content)) !== null) {
regions.push([match.index, match.index + match[0].length]);
@ -51,9 +56,38 @@ function findCodeBlockRegions(content: string): Array<[number, number]> {
}
/**
* Binary search to check if a position falls inside any code region.
* Match an inline link/image `[text](DEST)`, capturing the destination as group 1
* with the `d` flag so its span is read straight from `match.indices` (the text
* can contain an escaped `\](`, so a string search for the separator is unsafe).
* The text disallows unescaped `]`; the destination allows escapes and one level
* of balanced parens.
*/
function isInCodeBlock(
const LINK_DEST_RE =
/!?\[(?:\\.|[^\]\\])*?\]\(((?:\\.|[^()\\]|\([^()]*\))*)\)/gd;
/**
* Find the destination spans of inline links/images, so a `\(...\)` written with
* escaped parens inside a URL isn't rewritten as math (which would break the
* link). Only the destination is returned, not the link text, so math in the
* visible text still converts. Sorted, non-overlapping (matches are disjoint).
*/
function findLinkDestinationRegions(content: string): Array<[number, number]> {
if (!content.includes("](")) return [];
const regions: Array<[number, number]> = [];
let match: RegExpExecArray | null;
LINK_DEST_RE.lastIndex = 0;
while ((match = LINK_DEST_RE.exec(content)) !== null) {
// `indices` is present (the `d` flag); group 1 spans the destination.
regions.push(match.indices![1]);
}
return regions;
}
/**
* Binary search to check if a position falls inside any region. Regions must be
* sorted by start and non-overlapping.
*/
function isInRegion(
position: number,
regions: Array<[number, number]>,
): boolean {
@ -174,9 +208,109 @@ function hasInlineMathCloser(content: string, offset: number): boolean {
}
/**
* Preprocess a markdown string to escape currency dollar signs so they are not
* parsed as LaTeX math delimiters.
* Matches a `\[...\]` (display) or `\(...\)` (inline) LaTeX span. Non-greedy so
* the first closer wins; dotall so display spans can wrap lines. `(?<!\\)` on
* each opener leaves an escaped literal `\\[` (a real backslash then bracket)
* alone. The body is length-capped so an unclosed opener can't scan to
* end-of-input: without the cap, many unclosed `\(`/`\[` make matching O(n^2),
* and this runs per animation frame while streaming. Real spans are far shorter
* than the cap; a longer one just stays literal.
*/
const LATEX_DELIM_RE =
/(?<!\\)\\\[([\s\S]{0,4096}?)\\\]|(?<!\\)\\\(([\s\S]{0,4096}?)\\\)/g;
/**
* Rewrite `\[...\]` -> block `$$...$$` and `\(...\)` -> inline `$...$` so
* remark-math can tokenize them. Bodies are trimmed: remark-math won't open an
* inline span on `$ ` (a `$` followed by whitespace), and display fences must
* sit on their own line to render as a centered block (not inline math), so
* `\[...\]` becomes `\n$$\n...\n$$\n`.
*
* Spans inside code blocks/spans are left intact (a code sample showing `\(x\)`
* must not be rewritten).
*
* A space is inserted between a converted span and a following `$` so their
* delimiters can't fuse (`\(a\)\(b\)` -> `$a$$b$` would mis-tokenize into one
* broken span). A preceding currency (`$5\(x\)`) is instead broken later by the
* currency escape pass.
*
* Returns the rewritten text and the `[start, end)` ranges (in the rewritten
* string) of every span it produced, so the currency pass can skip them.
*/
function convertLatexDelimiters(content: string): {
text: string;
mathRegions: Array<[number, number]>;
} {
if (!content.includes("\\[") && !content.includes("\\(")) {
return { text: content, mathRegions: [] };
}
const codeRegions = findCodeBlockRegions(content);
const linkRegions = findLinkDestinationRegions(content);
const inSkipZone = (pos: number) =>
isInRegion(pos, codeRegions) || isInRegion(pos, linkRegions);
// Pushed in ascending, non-overlapping order (offset only grows), so this
// stays valid for isInRegion's binary search without a sort.
const mathRegions: Array<[number, number]> = [];
// Accumulate into an array, not a string: reading the last char off a growing
// `+=` accumulator flattens its rope every append (O(n^2) over many spans, on
// the per-frame streaming path), so track the tail char and length instead.
const parts: string[] = [];
let offset = 0;
let lastChar = "";
let last = 0;
// Append a chunk, separating a trailing `$` from a leading `$` so two spans
// can't fuse. Returns where the chunk landed (after any inserted space).
const append = (chunk: string): number => {
if (!chunk) return offset;
if (lastChar === "$" && chunk.startsWith("$")) {
parts.push(" ");
offset += 1;
}
const start = offset;
parts.push(chunk);
offset += chunk.length;
lastChar = chunk[chunk.length - 1];
return start;
};
let match: RegExpExecArray | null;
LATEX_DELIM_RE.lastIndex = 0;
while ((match = LATEX_DELIM_RE.exec(content)) !== null) {
const matchEnd = match.index + match[0].length;
// Skip if either delimiter is inside code or a link destination: an opener
// outside such a zone must not consume a closer inside one and rewrite
// across the boundary. Resume right after this opener (not past the whole
// match) so a valid span that this match spanned across (a stray code `\(`
// paired with a real closer) is still found on the next pass, not swallowed.
if (inSkipZone(match.index) || inSkipZone(matchEnd - 1)) {
LATEX_DELIM_RE.lastIndex = match.index + 1;
continue;
}
const isDisplay = match[1] !== undefined;
const body = (isDisplay ? match[1] : match[2]).trim();
// Leave an empty span (`\(\)`) literal; a bare `$$` would open a stray
// display block that swallows following text.
if (!body) {
continue;
}
append(content.slice(last, match.index));
const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`;
const start = append(wrapped);
mathRegions.push([start, offset]);
last = matchEnd;
}
append(content.slice(last));
return { text: parts.join(""), mathRegions };
}
/**
* Preprocess a markdown string so LaTeX renders: convert bracket delimiters to
* dollar forms, then escape currency dollar signs so they are not parsed as
* math delimiters.
*
* - `\[E = mc^2\]` becomes a `$$` display block on its own lines (display math)
* - `\(\alpha\)` becomes `$\alpha$` (inline math)
* - `\(x\)` in a code span is untouched
* - `$5` alone becomes `\$5` (currency, not math)
* - `$\alpha$` is untouched (real LaTeX)
* - `$30^\circ$` is untouched (LaTeX whose body starts with a digit)
@ -185,15 +319,22 @@ function hasInlineMathCloser(content: string, offset: number): boolean {
* - Currency inside code blocks/spans is untouched
*/
export function preprocessLaTeX(content: string): string {
if (!content.includes("$")) return content;
const { text, mathRegions } = convertLatexDelimiters(content);
const codeRegions = findCodeBlockRegions(content);
if (!text.includes("$")) return text;
return content.replace(CURRENCY_REGEX, (match, offset) => {
if (isInCodeBlock(offset, codeRegions)) {
const codeRegions = findCodeBlockRegions(text);
return text.replace(CURRENCY_REGEX, (match, offset) => {
if (isInRegion(offset, codeRegions)) {
return match;
}
if (hasInlineMathCloser(content, offset)) {
// Skip the spans we just created from `\(...\)` so a numeric body like
// `$5$` isn't re-escaped back to literal `\$5$`.
if (isInRegion(offset, mathRegions)) {
return match;
}
if (hasInlineMathCloser(text, offset)) {
return match;
}
return "\\" + match;

16
studio/frontend/src/main.tsx
View file

@ -5,8 +5,8 @@ import { StrictMode } from "react";
import { createRoot } from "react-dom/client";
import "./index.css";
import { fetchDeviceType } from "./config/env";
import { App } from "./app/app";
import { fetchDeviceType } from "./config/env";
import { initializeLocale } from "./i18n";
const globalCrypto = globalThis.crypto as Crypto | undefined;
@ -36,10 +36,10 @@ if (!rootElement) {
initializeLocale();
fetchDeviceType().then(() => {
createRoot(rootElement).render(
<StrictMode>
<App />
</StrictMode>,
);
});
createRoot(rootElement).render(
<StrictMode>
<App />
</StrictMode>,
);
fetchDeviceType().catch(() => undefined);

132
studio/src-tauri/src/preflight.rs
View file

@ -498,19 +498,68 @@ mod tests {
}
#[cfg(unix)]
fn remove_managed_capability_cache() {
let _ = std::fs::remove_file(
dirs::home_dir()
static MANAGED_CAPABILITY_CACHE_TEST_LOCK: std::sync::LazyLock<tokio::sync::Mutex<()>> =
std::sync::LazyLock::new(|| tokio::sync::Mutex::new(()));
#[cfg(unix)]
struct ManagedCapabilityCacheHome {
path: PathBuf,
previous: Option<std::ffi::OsString>,
}
#[cfg(unix)]
impl ManagedCapabilityCacheHome {
fn new(test_name: &str) -> Self {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.join(".unsloth")
.join("studio")
.join("desktop_capability_cache.json"),
);
.as_nanos();
let path = std::env::temp_dir().join(format!(
"unsloth-preflight-cache-{test_name}-{}-{nanos}",
std::process::id()
));
std::fs::create_dir_all(&path).unwrap();
let previous = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME");
std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", &path);
Self { path, previous }
}
}
#[cfg(unix)]
impl Drop for ManagedCapabilityCacheHome {
fn drop(&mut self) {
if let Some(previous) = &self.previous {
std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", previous);
} else {
std::env::remove_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME");
}
let _ = std::fs::remove_dir_all(&self.path);
}
}
#[cfg(unix)]
fn managed_capability_cache_path_for_test() -> PathBuf {
std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME")
.map(PathBuf::from)
.or_else(dirs::home_dir)
.unwrap()
.join(".unsloth")
.join("studio")
.join("desktop_capability_cache.json")
}
#[cfg(unix)]
fn remove_managed_capability_cache() {
let _ = std::fs::remove_file(managed_capability_cache_path_for_test());
}
#[cfg(unix)]
#[tokio::test]
async fn managed_cli_capability_probe_classifies_core_cases() {
let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await;
let _cache_home = ManagedCapabilityCacheHome::new("core-cases");
remove_managed_capability_cache();
for (name, script, stale_reason) in [
@ -567,6 +616,75 @@ exit 1
}
}
#[cfg(unix)]
#[tokio::test]
async fn managed_cli_capability_help_probe_runs_before_cache() {
use std::fs;
let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await;
let _cache_home = ManagedCapabilityCacheHome::new("cache-hit");
remove_managed_capability_cache();
// `-h` always succeeds unless `modeh` exists; the desktop-capabilities
// probe always succeeds unless `modecap` exists. Toggling those lets us
// prove the ordering: -h runs on every probe (even a cache hit), while
// the heavier capability probe is skipped once the cache is warm.
let fake = fake_cli(
"cap-cache-hit",
r#"#!/bin/sh
log="$0.calls"
modeh="$0.modeh"
modecap="$0.modecap"
printf '%s\n' "$*" >> "$log"
if [ "$1" = "-h" ]; then
if [ -f "$modeh" ]; then exit 42; fi
exit 0
fi
if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then
if [ -f "$modecap" ]; then exit 42; fi
printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}'
exit 0
fi
exit 1
"#,
);
let bin = fake.bin.clone();
let calls = bin.with_extension("calls");
let modeh = bin.with_extension("modeh");
let modecap = bin.with_extension("modecap");
// Cold probe: runs -h and the capability probe, then caches the result.
assert!(matches!(
probe_managed_bin(bin.clone()).await,
ManagedProbe::Ready { .. }
));
let first_calls = fs::read_to_string(&calls).unwrap();
assert!(first_calls.contains("-h"));
assert!(first_calls.contains("studio desktop-capabilities --json"));
// Cache hit: -h still runs, but the capability probe is skipped (breaking
// it via `modecap` proves it is not invoked).
fs::write(&modecap, "broken").unwrap();
fs::write(&calls, "").unwrap();
assert!(matches!(
probe_managed_bin(bin.clone()).await,
ManagedProbe::Ready { .. }
));
assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n");
// A non-launchable CLI is caught by the -h probe even with a warm cache:
// preflight reports Stale (for repair) and never trusts the cache.
fs::write(&modeh, "broken").unwrap();
fs::write(&calls, "").unwrap();
assert!(matches!(
probe_managed_bin(bin).await,
ManagedProbe::Stale { .. }
));
assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n");
remove_managed_capability_cache();
}
const EXPECTED_ROOT_ID: &str =
"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
const OTHER_ROOT_ID: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb";

16
studio/src-tauri/src/preflight/managed.rs
View file

@ -188,6 +188,16 @@ fn managed_bin_fingerprint(bin: &Path) -> Option<ManagedBinFingerprint> {
}
fn capability_cache_path() -> Option<PathBuf> {
#[cfg(test)]
if let Some(home) = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") {
return Some(
PathBuf::from(home)
.join(".unsloth")
.join("studio")
.join("desktop_capability_cache.json"),
);
}
dirs::home_dir().map(|home| {
home.join(".unsloth")
.join("studio")
@ -400,6 +410,12 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool {
pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
let started = Instant::now();
// Always verify the managed CLI actually launches before trusting the cache.
// A matching capability fingerprint does not prove the binary can still run:
// its venv interpreter or a runtime dependency can be broken while the
// path/size/mtime/markers are unchanged, so the -h probe runs first and a
// non-launchable install is reported Stale for repair. The capability cache
// below still skips the heavier desktop-capabilities probe on a hit.
if !run_cli_probe(&bin, &["-h"]).await {
info!(
"Managed preflight: cli unusable for {:?} in {}ms",

84
tests/_zoo_rocm_spoof.py Normal file
View file

@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""ROCm/RDNA spoof: present torch as an AMD Radeon (RDNA 2/3/4) card on a
GPU-less host, so hip paths (device_type -> "hip", llama.cpp ROCm bundle) are
testable in CPU-only CI with no AMD hardware. The ROCm sibling of
_zoo_aggressive_cuda_spoof.py: it reuses that spoof's torch.cuda no-op machinery
and overlays the AMD identity (torch.version.hip, gcnArchName, Radeon name).
Apply BEFORE importing unsloth/unsloth_zoo, since DEVICE_TYPE is cached there.
"""
from __future__ import annotations
import importlib.util
import os
import sys
# gfx -> (marketing name, (capability major, minor), torch.version.hip). hip is
# the ROCm build torch was made against (RDNA2/3 ship 6.x; gfx1102/115x/RDNA4 7.2).
_PROFILES: dict[str, tuple[str, tuple[int, int], str]] = {
"gfx1030": ("AMD Radeon RX 6900 XT", (10, 3), "6.4.43483"), # RDNA2
"gfx1031": ("AMD Radeon RX 6700 XT", (10, 3), "6.4.43483"),
"gfx1032": ("AMD Radeon RX 6600", (10, 3), "6.4.43483"),
"gfx1034": ("AMD Radeon RX 6400", (10, 3), "6.4.43483"),
"gfx1100": ("AMD Radeon RX 7900 XTX", (11, 0), "6.4.43483"), # RDNA3
"gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"),
"gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"),
"gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU
"gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"),
"gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4
"gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"),
}
def _cuda_spoof():
"""Load the sibling CUDA spoof by path (robust to sys.path), so we reuse its
torch.cuda machinery instead of duplicating it."""
if "_zoo_aggressive_cuda_spoof" in sys.modules:
return sys.modules["_zoo_aggressive_cuda_spoof"]
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py")
spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
sys.modules["_zoo_aggressive_cuda_spoof"] = mod
return mod
def apply(gfx: str = "gfx1100", device_count: int = 1) -> None:
"""Present torch as `gfx`. Re-callable to switch arch (identity is overlaid;
the underlying no-op machinery is applied once)."""
import torch
if gfx not in _PROFILES:
raise KeyError(f"Unknown gfx {gfx!r}; known: {', '.join(_PROFILES)}")
name, cap, hip = _PROFILES[gfx]
_cuda_spoof().apply() # is_available/device_count/streams/rng/amp/...
# Overlay the AMD identity on top of the (NVIDIA-shaped) CUDA spoof.
torch.version.hip = hip
torch.version.cuda = None
torch.cuda.device_count = lambda: device_count
torch.cuda.get_device_name = lambda *a, **k: name
torch.cuda.get_device_capability = lambda *a, **k: cap
torch.cuda.get_arch_list = lambda: [gfx]
class _Props:
pass
_p = _Props()
_p.name = name
_p.gcnArchName = f"{gfx}:sramecc-:xnack-" # ROCm advertises feature flags
_p.major, _p.minor = cap
_p.total_memory = 16 * 1024**3
_p.multi_processor_count = 40
_p.warp_size = 32 # RDNA wavefront (CDNA is 64)
_p.is_integrated = gfx in ("gfx1150", "gfx1151")
_p.is_multi_gpu_board = False
torch.cuda.get_device_properties = lambda *a, **k: _p
if __name__ == "__main__":
apply()
import torch
print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName)

459
tests/saving/test_prewarm_base_model_hub_cache.py Normal file
View file

@ -0,0 +1,459 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Regression tests for #6890: repeated base-model downloads across checkpoint exports.
merge_and_overwrite_lora downloads missing 16-bit shards with hf_hub_download(local_dir),
which never populates the persistent HF hub cache; a temporary merge directory (Studio
GGUF exports delete it) means every checkpoint export re-downloads the full base model.
_prewarm_base_model_hub_cache snapshot-downloads the base into the hub cache first so
the zoo's cache-copy fast path is hit on later exports.
unsloth.save cannot be imported on GPU-less hosts, so these tests extract the helper's
source via ast and exec it against fakes, mirroring the other GPU-free tests.
"""
from __future__ import annotations
import ast
import json
import os
import types
from pathlib import Path
import pytest
_SAVE_PY = Path(__file__).resolve().parent.parent.parent / "unsloth" / "save.py"
_SOURCE = _SAVE_PY.read_text(encoding = "utf-8")
def _extract_function(name: str) -> str:
tree = ast.parse(_SOURCE)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == name:
return ast.get_source_segment(_SOURCE, node)
raise AssertionError(f"{name} not found in unsloth/save.py")
class _FakePeftModel:
def __init__(self, name_or_path = "unsloth/gemma-4-31b-it-bnb-4bit"):
self.config = types.SimpleNamespace(_name_or_path = name_or_path)
class _Recorder:
"""Callable that records calls and returns/raises per configuration."""
def __init__(
self,
result = None,
exc = None,
results_fn = None,
):
self.calls = []
self.result = result
self.exc = exc
self.results_fn = results_fn
def __call__(self, *args, **kwargs):
self.calls.append((args, kwargs))
if self.exc is not None:
raise self.exc
if self.results_fn is not None:
return self.results_fn(*args, **kwargs)
return self.result
def _build_env(
monkeypatch,
tmp_path,
shards = None,
cached = False,
free_bytes = 10**15,
base_source = None,
kaggle = False,
colab = False,
hub_cache = None,
live_hub_cache = "__same__",
fp8_sibling = None,
sibling_source = None,
index_weight_map = None,
):
"""Exec the extracted helper with stubbed collaborators; returns (fn, stubs)."""
shards = (
shards
if shards is not None
else [
("model-00001-of-00002.safetensors", 30 * 1024**3),
("model-00002-of-00002.safetensors", 29 * 1024**3),
]
)
if base_source is None:
base_source = ("unsloth/gemma-4-31b-it", False, None, False, None)
class _FS:
def __init__(self, token = None):
pass
def ls(
self,
repo,
detail = True,
):
return [{"name": f"{repo}/{n}", "size": s} for n, s in shards]
class _LocalMiss(Exception):
pass
hf_hub_download = _Recorder(result = str(tmp_path / "cached"))
if not cached:
hf_hub_download.exc = _LocalMiss("not cached")
# Serve model.safetensors.index.json (the merge's shard filter) while shard cache
# probes still miss, so the index-filter path can be exercised without a network.
if index_weight_map is not None:
_idx_path = tmp_path / "model.safetensors.index.json"
_idx_path.write_text(json.dumps({"weight_map": index_weight_map}))
def _hub_dl(
repo_id = None,
filename = None,
**kw,
):
if filename == "model.safetensors.index.json":
return str(_idx_path)
raise _LocalMiss("not cached")
hf_hub_download.exc = None
hf_hub_download.results_fn = _hub_dl
snapshot_download = _Recorder()
determine_base_model_source = _Recorder(result = base_source)
# For the FP8 -> 16bit sibling swap: return the sibling's (16bit) source when the
# helper re-resolves the sibling, else the original base source.
if fp8_sibling is not None:
_sib_src = sibling_source or (fp8_sibling, False, None, False, None)
determine_base_model_source.results_fn = (
lambda name, token = None: _sib_src if name == fp8_sibling else base_source
)
resolve_fp8_16bit_sibling = _Recorder(result = fp8_sibling)
cache_dir = tmp_path / "hub_cache"
cache_dir.mkdir(exist_ok = True)
hf_module = types.SimpleNamespace(
HfFileSystem = _FS,
hf_hub_download = hf_hub_download,
snapshot_download = snapshot_download,
constants = types.SimpleNamespace(
HF_HUB_CACHE = hub_cache if hub_cache is not None else str(cache_dir)
),
)
zoo_module = types.SimpleNamespace(
determine_base_model_source = determine_base_model_source,
_resolve_fp8_16bit_sibling = resolve_fp8_16bit_sibling,
)
# Stub the live-env cache resolver the pre-warm uses (matches what the merge reads).
_live = (
(hub_cache if hub_cache is not None else str(cache_dir))
if live_hub_cache == "__same__"
else live_hub_cache
)
hf_cache_module = types.SimpleNamespace(_active_caches = lambda: (None, _live, None))
monkeypatch.setitem(__import__("sys").modules, "huggingface_hub", hf_module)
monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.saving_utils", zoo_module)
monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.hf_cache", hf_cache_module)
fake_shutil = types.SimpleNamespace(
disk_usage = lambda path: types.SimpleNamespace(free = free_bytes)
)
prints = []
namespace = {
"os": os,
"shutil": fake_shutil,
"PeftModel": _FakePeftModel,
"get_model_name": lambda name, load_in_4bit: name.removesuffix("-bnb-4bit"),
"IS_KAGGLE_ENVIRONMENT": kaggle,
"IS_COLAB_ENVIRONMENT": colab,
"print": lambda *a, **k: prints.append(" ".join(str(x) for x in a)),
}
exec(
compile(_extract_function("_prewarm_base_model_hub_cache"), str(_SAVE_PY), "exec"),
namespace,
)
stubs = types.SimpleNamespace(
snapshot_download = snapshot_download,
hf_hub_download = hf_hub_download,
determine_base_model_source = determine_base_model_source,
resolve_fp8_16bit_sibling = resolve_fp8_16bit_sibling,
prints = prints,
)
return namespace["_prewarm_base_model_hub_cache"], stubs
def test_downloads_base_into_hub_cache(monkeypatch, tmp_path):
monkeypatch.delenv("UNSLOTH_PREWARM_HUB_CACHE", raising = False)
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(_FakePeftModel(), save_method = "merged_16bit", token = "tok")
assert len(stubs.snapshot_download.calls) == 1
_, kwargs = stubs.snapshot_download.calls[0]
assert kwargs["repo_id"] == "unsloth/gemma-4-31b-it"
# No local_dir: the whole point is populating the persistent cache.
assert "local_dir" not in kwargs
assert "model-00001-of-00002.safetensors" in kwargs["allow_patterns"]
assert "model.safetensors.index.json" in kwargs["allow_patterns"]
def test_skips_download_when_already_cached(monkeypatch, tmp_path):
fn, stubs = _build_env(monkeypatch, tmp_path, cached = True)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == []
# The cached check must not hit the network.
assert all(kwargs.get("local_files_only") for _, kwargs in stubs.hf_hub_download.calls)
def test_skips_when_disk_too_small_for_cache_copy(monkeypatch, tmp_path):
fn, stubs = _build_env(monkeypatch, tmp_path, free_bytes = 60 * 1024**3)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == []
@pytest.mark.parametrize("env_value", ["0", "false", "NO", "off"])
def test_env_opt_out(monkeypatch, tmp_path, env_value):
monkeypatch.setenv("UNSLOTH_PREWARM_HUB_CACHE", env_value)
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == []
@pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"])
def test_offline_skips(monkeypatch, tmp_path, var):
monkeypatch.setenv(var, "1")
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == []
def test_kaggle_and_colab_skip(monkeypatch, tmp_path):
for flag in ("kaggle", "colab"):
fn, stubs = _build_env(monkeypatch, tmp_path, **{flag: True})
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == [], f"{flag} must skip pre-warm"
@pytest.mark.parametrize("save_method", ["merged_4bit", "forced_merged_4bit", "lora"])
def test_non_downloading_save_methods_skip(monkeypatch, tmp_path, save_method):
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(_FakePeftModel(), save_method = save_method)
assert stubs.snapshot_download.calls == []
def test_local_base_model_skips(monkeypatch, tmp_path):
local_dir = tmp_path / "local_base"
local_dir.mkdir()
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(_FakePeftModel(name_or_path = str(local_dir)), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == []
def test_quantized_base_skips(monkeypatch, tmp_path):
fn, stubs = _build_env(
monkeypatch,
tmp_path,
base_source = ("unsloth/gemma-4-31b-it-bnb-4bit", False, None, True, "nf4"),
)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls == []
def test_non_peft_model_skips(monkeypatch, tmp_path):
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(object(), save_method = "merged_16bit")
assert stubs.determine_base_model_source.calls == []
assert stubs.snapshot_download.calls == []
def test_consolidated_shard_excluded_when_proper_shards_exist(monkeypatch, tmp_path):
fn, stubs = _build_env(
monkeypatch,
tmp_path,
shards = [
("consolidated.safetensors", 14 * 1024**3),
("model-00001-of-00001.safetensors", 14 * 1024**3),
],
)
fn(_FakePeftModel(), save_method = "merged_16bit")
_, kwargs = stubs.snapshot_download.calls[0]
assert "consolidated.safetensors" not in kwargs["allow_patterns"]
assert "model-00001-of-00001.safetensors" in kwargs["allow_patterns"]
def test_consolidated_only_repo_is_kept(monkeypatch, tmp_path):
fn, stubs = _build_env(
monkeypatch,
tmp_path,
shards = [("consolidated.safetensors", 14 * 1024**3)],
)
fn(_FakePeftModel(), save_method = "merged_16bit")
_, kwargs = stubs.snapshot_download.calls[0]
assert "consolidated.safetensors" in kwargs["allow_patterns"]
def test_listing_failure_is_swallowed(monkeypatch, tmp_path):
fn, stubs = _build_env(monkeypatch, tmp_path)
stubs.determine_base_model_source.exc = RuntimeError("HF is down")
fn(_FakePeftModel(), save_method = "merged_16bit") # must not raise
assert stubs.snapshot_download.calls == []
def test_gpt_oss_bf16_mxfp4_swap_skips(monkeypatch, tmp_path):
fn, stubs = _build_env(monkeypatch, tmp_path)
fn(
_FakePeftModel(name_or_path = "unsloth/gpt-oss-20b-BF16"),
save_method = "mxfp4",
)
assert stubs.snapshot_download.calls == []
def test_missing_config_skips_cleanly(monkeypatch, tmp_path):
# A model whose config is None must skip silently, not fall into the outer
# exception handler that prints a misleading "Could not pre-cache" warning.
fn, stubs = _build_env(monkeypatch, tmp_path)
model = _FakePeftModel() # a PeftModel instance so the isinstance guard passes
model.config = None
fn(model, save_method = "merged_16bit")
assert stubs.determine_base_model_source.calls == []
assert stubs.snapshot_download.calls == []
assert not any(
"Could not pre-cache" in p for p in stubs.prints
), "missing config took the error path instead of a clean skip"
def test_relative_hub_cache_does_not_falsely_skip(monkeypatch, tmp_path):
# A relative HF_HUB_CACHE whose leaf does not exist yet must still resolve to a real
# root for the disk probe; without abspath the walk-up hits "" and pre-warm is skipped.
monkeypatch.chdir(tmp_path)
fn, stubs = _build_env(monkeypatch, tmp_path, hub_cache = "relcache/hub")
fn(_FakePeftModel(), save_method = "merged_16bit")
assert len(stubs.snapshot_download.calls) == 1, "relative cache path falsely skipped pre-warm"
def test_generic_save_calls_prewarm_before_merge():
"""unsloth_generic_save must pre-warm the cache before merge_and_overwrite_lora."""
tree = ast.parse(_SOURCE)
fn = next(
node
for node in ast.walk(tree)
if isinstance(node, ast.FunctionDef) and node.name == "unsloth_generic_save"
)
body_src = ast.get_source_segment(_SOURCE, fn)
prewarm_pos = body_src.find("_prewarm_base_model_hub_cache(")
merge_pos = body_src.find("merge_and_overwrite_lora(")
assert prewarm_pos != -1, "unsloth_generic_save no longer pre-warms the hub cache"
assert merge_pos != -1
assert prewarm_pos < merge_pos, "pre-warm must run before the merge downloads shards"
def test_prewarm_downloads_into_live_env_cache(monkeypatch, tmp_path):
# Download must target the live-env cache (what the merge reads), via cache_dir.
fn, stubs = _build_env(monkeypatch, tmp_path, live_hub_cache = "/mnt/persistent/hf/hub")
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/mnt/persistent/hf/hub"
def test_prewarm_survives_runtime_cache_redirect(monkeypatch, tmp_path):
# Frozen constants (stale dir) vs the merge's runtime-redirected dir: the pre-warm
# must follow the redirect, else the cache-copy fast path misses and #6890 is unfixed.
fn, stubs = _build_env(
monkeypatch,
tmp_path,
hub_cache = "/read-only/original/hub",
live_hub_cache = "/writable/redirect/hub",
)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/writable/redirect/hub"
def test_cached_probe_uses_live_env_cache(monkeypatch, tmp_path):
# The already-cached fast path must probe the live-env cache dir too.
fn, stubs = _build_env(
monkeypatch, tmp_path, cached = True, live_hub_cache = "/mnt/persistent/hf/hub"
)
fn(_FakePeftModel(), save_method = "merged_16bit")
assert stubs.hf_hub_download.calls, "cached probe did not run"
assert all(
kw.get("cache_dir") == "/mnt/persistent/hf/hub" for _, kw in stubs.hf_hub_download.calls
)
def test_fp8_base_prewarms_16bit_sibling_not_fp8_repo(monkeypatch, tmp_path):
# A merged_16bit export of an FP8 base with a 16bit sibling merges onto the sibling,
# so the pre-warm must cache the sibling (what the merge downloads), not the FP8 repo.
fn, stubs = _build_env(
monkeypatch,
tmp_path,
base_source = ("unsloth/Model-FP8", False, None, True, "fp8"),
fp8_sibling = "unsloth/Model",
)
fn(_FakePeftModel(name_or_path = "unsloth/Model-FP8"), save_method = "merged_16bit")
assert stubs.resolve_fp8_16bit_sibling.calls, "sibling resolver was not consulted"
assert len(stubs.snapshot_download.calls) == 1
assert stubs.snapshot_download.calls[0][1]["repo_id"] == "unsloth/Model"
def test_fp8_base_without_sibling_still_prewarms_fp8_repo(monkeypatch, tmp_path):
# No sibling: the merge dequants the FP8 base in place, so caching the FP8 repo helps.
fn, stubs = _build_env(
monkeypatch,
tmp_path,
base_source = ("unsloth/Model-FP8", False, None, True, "fp8"),
fp8_sibling = None,
)
fn(_FakePeftModel(name_or_path = "unsloth/Model-FP8"), save_method = "merged_16bit")
assert len(stubs.snapshot_download.calls) == 1
assert stubs.snapshot_download.calls[0][1]["repo_id"] == "unsloth/Model-FP8"
def test_prewarm_filters_shards_through_index(monkeypatch, tmp_path):
# A repo with a leftover shard not referenced by the index: the merge keeps only the
# indexed shards, so the pre-warm must too (else the disk gate over-counts and
# snapshot_download fetches the unused leftover).
fn, stubs = _build_env(
monkeypatch,
tmp_path,
shards = [
("model-00001-of-00002.safetensors", 10 * 1024**3),
("model-00002-of-00002.safetensors", 10 * 1024**3),
("leftover-00001-of-00001.safetensors", 10 * 1024**3),
],
index_weight_map = {
"a.weight": "model-00001-of-00002.safetensors",
"b.weight": "model-00002-of-00002.safetensors",
},
)
fn(_FakePeftModel(), save_method = "merged_16bit")
allow = stubs.snapshot_download.calls[0][1]["allow_patterns"]
assert "leftover-00001-of-00001.safetensors" not in allow
assert "model-00001-of-00002.safetensors" in allow
assert "model-00002-of-00002.safetensors" in allow
def test_prewarm_keeps_all_shards_when_index_matches(monkeypatch, tmp_path):
# No leftover: every listed shard is indexed, so none are dropped.
fn, stubs = _build_env(
monkeypatch,
tmp_path,
shards = [
("model-00001-of-00002.safetensors", 10 * 1024**3),
("model-00002-of-00002.safetensors", 10 * 1024**3),
],
index_weight_map = {
"a.weight": "model-00001-of-00002.safetensors",
"b.weight": "model-00002-of-00002.safetensors",
},
)
fn(_FakePeftModel(), save_method = "merged_16bit")
allow = stubs.snapshot_download.calls[0][1]["allow_patterns"]
assert "model-00001-of-00002.safetensors" in allow
assert "model-00002-of-00002.safetensors" in allow

84
tests/studio/install/test_rocm_rdna_routing.py Normal file
View file

@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""RDNA 2/3/4 routing, validated on CPU-only CI with no AMD hardware.
tests/_zoo_rocm_spoof.py presents torch as each Radeon gfx arch, then we assert
unsloth_zoo routes it: device_type -> "hip", llama.cpp target -> ("rocm", gfx),
and the per-family ROCm bundle suffix. The torch-facing checks run in a
subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached
at import) resolves from a clean process.
"""
from __future__ import annotations
import json
import subprocess
import sys
from pathlib import Path
import pytest
pytest.importorskip("torch")
pytest.importorskip("unsloth_zoo")
_TESTS_DIR = Path(__file__).resolve().parents[2] # tests/
# gfx -> (expected llama.cpp target, expected ROCm bundle family).
_ARCHES = {
"gfx1030": (("rocm", "gfx1030"), "gfx103X"), # RDNA2
"gfx1031": (("rocm", "gfx1031"), "gfx103X"),
"gfx1032": (("rocm", "gfx1032"), "gfx103X"),
"gfx1034": (("rocm", "gfx1034"), "gfx103X"),
"gfx1100": (("rocm", "gfx1100"), "gfx110X"), # RDNA3
"gfx1101": (("rocm", "gfx1101"), "gfx110X"),
"gfx1102": (("rocm", "gfx1102"), "gfx110X"),
"gfx1150": (("rocm", "gfx1150"), "gfx1150"), # RDNA3.5 APU (self-family)
"gfx1151": (("rocm", "gfx1151"), "gfx1151"),
"gfx1200": (("rocm", "gfx1200"), "gfx120X"), # RDNA4
"gfx1201": (("rocm", "gfx1201"), "gfx120X"),
}
# Child: spoof each arch, then record device_type once (fresh import) and the
# live llama.cpp target per arch. Emits one JSON line the parent parses.
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
from unsloth_zoo.device_type import get_device_type, is_hip
device_type = [get_device_type(), is_hip()]
from unsloth_zoo import llama_cpp as lc
targets = {{}}
for gfx in arches:
spoof.apply(gfx)
targets[gfx] = list(lc._detect_gpu_target())
print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}}))
"""
@pytest.fixture(scope = "module")
def routed():
code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES))
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True)
line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None)
assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
return json.loads(line[len("RESULT ") :])
@pytest.mark.parametrize("gfx", list(_ARCHES))
def test_detect_gpu_target(routed, gfx):
# RDNA card is routed to its ROCm gfx target (drives the llama.cpp bundle).
assert tuple(routed["targets"][gfx]) == _ARCHES[gfx][0]
def test_device_type_is_hip(routed):
# An RDNA card must resolve the compute device_type to "hip".
assert routed["device_type"] == ["hip", True]
@pytest.mark.parametrize("gfx", list(_ARCHES))
def test_rocm_gfx_family(gfx):
# Pure mapping (no torch): each gfx picks the right per-family ROCm bundle.
from unsloth_zoo import llama_cpp as lc
assert lc._rocm_gfx_family(gfx) == _ARCHES[gfx][1]

48
tests/studio/install/test_rocm_support.py
View file

@ -3426,8 +3426,9 @@ class TestInstallShDropinPersistence:
def test_gate5_early_return_persists_dropin(self):
"""The rocminfo-already-works early return must call the persist helper before returning."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
# The persist call must precede `return 0` at the rocminfo gfx1151 gate.
gate = source.find("Name:[[:space:]]*gfx1151")
# The persist call must precede `return 0` at the rocminfo GPU-agent gate
# (uniquely identified by the `!/generic/` clause the other probes lack).
gate = source.find("Name:[[:space:]]*gfx[1-9]/ && !/generic/")
assert gate != -1
window = source[gate : gate + 900]
assert "_persist_rocm_wsl_dropin" in window
@ -3441,6 +3442,49 @@ class TestInstallShDropinPersistence:
assert "profile.d/unsloth-rocm-wsl.sh" in body
_STRIXHALO_WSL_PATH = PACKAGE_ROOT / "scripts" / "install_rocm_wsl_strixhalo.sh"
class TestWslRerouteNvidiaGuard:
"""_maybe_reroute_strixhalo_to_2404 must skip the AMD reroute on hybrid AMD+NVIDIA hosts by
reusing _has_usable_nvidia_gpu (CUDA_VISIBLE_DEVICES-aware + /proc/driver/nvidia fallback),
which must be defined before the reroute's call site so it is actually available."""
def test_reroute_calls_nvidia_helper_before_amd_signal(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
start = source.find("_maybe_reroute_strixhalo_to_2404()")
assert start != -1
body = source[start : start + 1200]
nv = body.find("_has_usable_nvidia_gpu")
wmi = body.find("_wsl_amd_gpu_name")
assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute"
assert wmi != -1
# The NVIDIA guard must precede the AMD/WMI signal and return early.
assert nv < wmi
assert body.find("return 0", nv) < wmi
def test_nvidia_helper_and_deps_defined_before_reroute_callsite(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
call = source.find("\n_maybe_reroute_strixhalo_to_2404 || true")
assert call != -1
for fn in ("_run_bounded() {", "_cvd_hides_nvidia() {", "_has_usable_nvidia_gpu() {"):
idx = source.find(fn)
assert idx != -1 and idx < call, f"{fn} must be defined before the reroute call"
class TestStrixhaloGfxOverridePipefail:
"""The UNSLOTH_WSL_GFX override check must use a consuming grep, not grep -q: under
`set -o pipefail` an early -q exit SIGPIPEs printf and misreports the arch on large output."""
def test_gfx_override_uses_consuming_grep(self):
source = _STRIXHALO_WSL_PATH.read_text(encoding = "utf-8")
idx = source.find('grep -E "Name:[[:space:]]*${GFX}')
assert idx != -1, "GFX override must use a consuming grep -E (not grep -q)"
line = source[idx : source.find("\n", idx)]
assert ">/dev/null" in line
assert 'grep -qE "Name:[[:space:]]*${GFX}' not in source
class TestLlamaCppRuntimeWslOrdering:
"""The serve-time launcher mirrors binary_env: system HIP before the bundle dir on WSL."""

93
tests/studio/test_chat_response_details_ui_contract.py Normal file
View file

@ -0,0 +1,93 @@
"""Static contract for the chat response-details action and metadata."""
from __future__ import annotations
import re
from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
THREAD_TSX = REPO / "studio/frontend/src/components/assistant-ui/thread.tsx"
DETAILS_TSX = (
REPO / "studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx"
)
REASONING_TSX = REPO / "studio/frontend/src/components/assistant-ui/reasoning.tsx"
ADAPTER_TS = REPO / "studio/frontend/src/features/chat/api/chat-adapter.ts"
CHAT_PREFS_TS = REPO / "studio/frontend/src/features/chat/stores/chat-preferences-store.ts"
CHAT_TAB_TSX = REPO / "studio/frontend/src/features/settings/tabs/chat-tab.tsx"
def test_assistant_more_menu_exposes_response_details_action():
src = THREAD_TSX.read_text()
assert "MessageResponseDetailsSheet" in src
assert "See response details" in src
assert "setDetailsOpen(true)" in src
def test_response_details_sheet_uses_unsloth_sheet_and_key_sections():
src = DETAILS_TSX.read_text()
assert "SheetContent" in src
assert "Response details" in src
assert "MessageResponseModelBadge" in src
assert "showResponseModel" in src
assert "ChipIcon" not in src
assert "s.params.checkpoint" not in src
assert "Not recorded" in src
assert "min-w-0 break-words font-heading" in src
assert "toolCallsFromContent(message.content)" in src
assert 'label="Called"' in src
for section in ["Response", "Tokens", "Timing", "Tools"]:
assert f'title="{section}"' in src
for field in ["Model", "Provider", "Total", "Cache hits", "Enabled", "Called"]:
assert f'label="{field}"' in src
def test_response_model_chip_is_user_configurable_and_rendered_in_metadata_rows():
prefs_src = CHAT_PREFS_TS.read_text()
chat_tab_src = CHAT_TAB_TSX.read_text()
thread_src = THREAD_TSX.read_text()
reasoning_src = REASONING_TSX.read_text()
assert "showResponseModel: boolean" in prefs_src
assert "showResponseModel: false" in prefs_src
assert "showResponseModel: saved?.showResponseModel ?? false" in prefs_src
assert "Show response model" in chat_tab_src
assert "setShowResponseModel" in chat_tab_src
assert "aui-response-model-badge inline-flex min-h-5" in DETAILS_TSX.read_text()
assert "leading-5" in DETAILS_TSX.read_text()
assert "group-hover/assistant-message:opacity-100" in DETAILS_TSX.read_text()
assert "MessageResponseModelBadge" in thread_src
assert "hasReasoningParts" in thread_src
assert "group/assistant-message aui-assistant-message-root" in thread_src
assert "pointer-events-none relative h-0" in thread_src
assert "MessageResponseModelBadge" in reasoning_src
assert 'className="min-w-0 flex-none"' in reasoning_src
assert "hidden min-w-0 max-w-[12rem]" in reasoning_src
assert "group-hover/assistant-message:inline-flex" in reasoning_src
def test_response_details_metadata_is_persisted_without_backend_schema_change():
src = ADAPTER_TS.read_text()
assert "interface ResponseDetailsMetadata" in src
assert "buildResponseDetails" in src
assert "responseDetails: buildResponseDetails(finishedAt)" in src
assert "toolCalls: Array.from(" in src
assert "!isExternalRequest && supportsTools && toolsEnabled" in src
assert "!isExternalRequest && supportsTools && codeToolsEnabled" in src
assert re.search(r"selectedModelSummary\?\.name\s*\|\|\s*responseModelId", src)
assert "providerName" in src
assert "cancelId" in src
metadata_block = src[
src.find("interface ResponseDetailsMetadata") : src.find("type RunMessages")
]
builder_block = src[
src.find("const buildResponseDetails") : src.find("const externalCapabilities")
]
for forbidden in [
"encrypted_api_key",
"externalApiKey",
"apiKey",
"providerKey",
"secret",
]:
assert forbidden not in metadata_block
assert forbidden not in builder_block

138
tests/utils/test_rope_scaling_drift.py
View file

@ -257,6 +257,63 @@ def test_recompute_helper_scales_on_cpu():
), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
def test_extended_rotary_reads_config_factor():
# LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8
# (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405).
from types import SimpleNamespace
from unsloth.models.llama import LlamaExtendedRotaryEmbedding
rot = object.__new__(LlamaExtendedRotaryEmbedding)
rot.base = ROPE_THETA
rot.dim = HEAD_DIM
rot._unsloth_rope_config = SimpleNamespace(
rope_scaling = {
"rope_type": "llama3",
"factor": 32.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
}
)
vanilla = _vanilla_inv_freq()
scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1)
ratio = float(vanilla[-1]) / float(scaled[-1])
assert abs(ratio - 32.0) < 1e-3, (
f"LlamaExtendedRotaryEmbedding ignored config factor 32 (ratio {ratio}); the "
"low-frequency band must be divided by the config factor (issue #2405)."
)
def test_extended_rotary_reads_rope_parameters_v5():
# transformers v5 stores scaling under rope_parameters (rope_scaling is a
# back-compat shim that may be removed); the factor must still be read.
from types import SimpleNamespace
from unsloth.models.llama import LlamaExtendedRotaryEmbedding
rot = object.__new__(LlamaExtendedRotaryEmbedding)
rot.base = ROPE_THETA
rot.dim = HEAD_DIM
rot._unsloth_rope_config = SimpleNamespace(
rope_scaling = None,
rope_parameters = {
"rope_type": "llama3",
"factor": 32.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
},
)
vanilla = _vanilla_inv_freq()
scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1)
ratio = float(vanilla[-1]) / float(scaled[-1])
assert abs(ratio - 32.0) < 1e-3, (
f"Extended rotary ignored rope_parameters factor 32 (ratio {ratio}); v5 "
"keeps the factor under rope_parameters, not rope_scaling."
)
def _cos_at_position(rot, position):
"""cos row at one position, built like _set_cos_sin_cache but CPU-only."""
inv_freq = rot.inv_freq.float().cpu()
@ -324,6 +381,87 @@ def test_extended_cache_keeps_scaling_after_growth():
)
def _blank_nonpersistent_buffers(module):
"""Mimic transformers v5 meta-load: overwrite non-persistent buffers with garbage."""
for name, buf in list(module.named_buffers()):
leaf = module
*parents, attr = name.split(".")
for part in parents:
leaf = getattr(leaf, part)
if attr in getattr(leaf, "_non_persistent_buffers_set", set()):
setattr(leaf, attr, torch.rand_like(buf))
def _build_llama3_rotary():
from unsloth.models import llama as llama_mod
config = _make_config(LLAMA3_ROPE_SCALING)
return llama_mod.LlamaRotaryEmbedding(config = config), config
def _build_longrope_rotary():
from types import SimpleNamespace
from unsloth.models import llama as llama_mod
short_factor, long_factor = [1.05] * 48, [1.3] * 48
rot = llama_mod.LongRopeRotaryEmbedding(
dim = 96,
max_position_embeddings = 131072,
original_max_position_embeddings = 4096,
base = ROPE_THETA,
short_factor = short_factor,
long_factor = long_factor,
)
config = SimpleNamespace(
rope_scaling = {
"rope_type": "longrope",
"short_factor": short_factor,
"long_factor": long_factor,
"original_max_position_embeddings": 4096,
}
)
return rot, config
@requires_cuda
@pytest.mark.parametrize(
"build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"]
)
def test_v5_blank_repair_roundtrip(build):
# Build scaled -> blank non-persistent buffers (what transformers v5 does on
# load) -> run the repair -> every buffer must return to its scaled value.
# Family-agnostic: encodes no scaling math, so it guards any rotary that
# keeps scaling in a buffer (issue #2405 / PR #6907).
from unsloth.models import loader
# The repair only runs on transformers v5 (it is what blanks the buffers);
# on v4 _fix_rope_inv_freq is a no-op, so the round-trip cannot restore.
if not loader._NEEDS_ROPE_FIX:
pytest.skip("transformers < 5 does not blank rope buffers; repair is a no-op")
rot, config = build()
snapshot = {name: buf.detach().clone() for name, buf in rot.named_buffers()}
assert snapshot, "rotary registers no buffers; nothing to guard"
_blank_nonpersistent_buffers(rot)
assert any(
not torch.equal(rot.get_buffer(name), snapshot[name]) for name in snapshot
), "blanking changed no buffer; the round-trip would be vacuous"
wrapper = torch.nn.Module()
wrapper.add_module("rotary_emb", rot)
wrapper.config = config
loader._fix_rope_inv_freq(wrapper)
for name in snapshot:
assert torch.allclose(
rot.get_buffer(name).cpu(), snapshot[name].cpu(), rtol = 1e-4, atol = 1e-6
), (
f"{name} was not restored to its scaled value by loader._fix_rope_inv_freq "
"after the transformers v5 buffer blank (issue #2405 / PR #6907)."
)
def test_object_style_rope_scaling_does_not_crash():
# Object-style rope_scaling must be normalized, not .get()'d directly.
from dataclasses import dataclass

3
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2026.6.9"
__version__ = "2026.7.1"
__all__ = [
"SUPPORTS_BFLOAT16",
@ -2834,6 +2834,7 @@ def patch_llama_rope_scaling(
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
base=self.rope_theta,
config=self.config,
)
elif scaling_type == "longrope":
self.rotary_emb = {longrope_rope_function}(

17
unsloth/models/llama.py
View file

@ -1930,11 +1930,18 @@ class LlamaExtendedRotaryEmbedding(LlamaRotaryEmbedding):
# From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41
def _apply_inv_freq_scaling(self, freqs: torch.Tensor):
# Values obtained from grid search
scale_factor = 8
low_freq_factor = 1
high_freq_factor = 4
old_context_len = 8192 # original llama3 length
# llama3 factors from config; Llama-3.1 defaults when built without one
# (legacy codegen path). Hardcoding 8 is wrong for e.g. Llama-3.2 (32).
# v5 renames rope_scaling -> rope_parameters; read either so the factor
# survives even if the rope_scaling back-compat shim is dropped.
config = getattr(self, "_unsloth_rope_config", None)
rope_scaling = _rope_scaling_as_dict(
getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {}
)
scale_factor = rope_scaling.get("factor", 8)
low_freq_factor = rope_scaling.get("low_freq_factor", 1)
high_freq_factor = rope_scaling.get("high_freq_factor", 4)
old_context_len = rope_scaling.get("original_max_position_embeddings", 8192)
low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor

3
unsloth/models/rl.py
View file

@ -1370,6 +1370,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
# [TODO] See https://fengyao.notion.site/off-policy-rl
# https://github.com/huggingface/trl/pull/3867 (August 7th)
"vllm_importance_sampling_correction": False,
# TRL >= 1.7.0 enables the MoE router aux loss by default (0.001); the optimized
# GRPO forward does not compute it, so default off. Opt in via router_aux_loss_coef > 0.
"router_aux_loss_coef": 0.0,
}
for k, v in replacements.items():
x = f"{k}( = [^,\n]{{1,}})?,\n"

177
unsloth/save.py
View file

@ -3693,6 +3693,182 @@ from unsloth_zoo.llama_cpp import (
)
def _prewarm_base_model_hub_cache(
model,
save_method = "merged_16bit",
token = None,
):
"""Download the 16-bit base weights into the persistent HF hub cache before the merge.
merge_and_overwrite_lora fetches missing shards with hf_hub_download(local_dir = ...),
which never populates the hub cache. When the merge directory is temporary (GGUF
checkpoint exports delete it after conversion), every export re-downloads the full
base model (#6890). Pre-warming the cache makes the first export download once and
later exports copy from the cache. Best-effort: any failure or skip falls back to
the streaming download. Disable with UNSLOTH_PREWARM_HUB_CACHE=0.
"""
_false = ("0", "false", "no", "off")
if os.environ.get("UNSLOTH_PREWARM_HUB_CACHE", "1").strip().lower() in _false:
return
if IS_KAGGLE_ENVIRONMENT or IS_COLAB_ENVIRONMENT:
return
_true = ("1", "true", "yes", "on")
if (
os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _true
or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _true
):
return
# Only the 16bit / mxfp4 merges download the base model; merged_4bit and lora do not.
if save_method not in ("merged_16bit", "mxfp4"):
return
if not isinstance(model, PeftModel):
return
try:
# getattr so a model without a config / _name_or_path skips instead of raising.
name_or_path = getattr(getattr(model, "config", None), "_name_or_path", None)
if not name_or_path:
return
try:
model_name = get_model_name(name_or_path, load_in_4bit = False)
except Exception:
model_name = name_or_path
if not model_name or os.path.isdir(model_name):
return # local checkpoints are copied, never downloaded
# The merge may swap a gpt-oss "-BF16" repo for its MXFP4 variant, so skip it.
if save_method == "mxfp4" and model_name.endswith("-BF16"):
return
from unsloth_zoo.saving_utils import determine_base_model_source
model_name, is_local_path, _, base_is_quantized, quant_type = determine_base_model_source(
model_name, token
)
if not model_name or is_local_path:
return
# Mirror the merge: an FP8 base with a 16bit sibling merges onto the sibling, so
# pre-warm the sibling (what the merge downloads), not the FP8 repo (#6890).
if base_is_quantized and quant_type == "fp8" and save_method == "merged_16bit":
try:
from unsloth_zoo.saving_utils import _resolve_fp8_16bit_sibling
sibling = _resolve_fp8_16bit_sibling(model_name, token)
except Exception:
sibling = None
if sibling:
model_name, is_local_path, _, base_is_quantized, quant_type = (
determine_base_model_source(sibling, token)
)
if not model_name or is_local_path:
return
if base_is_quantized and quant_type in ("nf4", "fp4"):
return # the 16bit merge refuses these bases; nothing worth caching
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
# Resolve the cache from the live env like the merge, not huggingface_hub's frozen
# constants: a runtime cache redirect (read-only default, Studio) would else miss (#6890).
try:
from unsloth_zoo.hf_cache import _active_caches
_hub_cache = _active_caches()[1]
hub_cache_dir = str(_hub_cache) if _hub_cache is not None else None
except Exception:
hub_cache_dir = None
# Mirror the zoo's shard listing (drop consolidated.safetensors when proper
# shards coexist) so the cached set is a superset of what the merge looks up.
shard_names = []
total_size_in_bytes = 0
for x in HfFileSystem(token = token).ls(model_name, detail = True):
if x["name"].endswith(".safetensors"):
shard_names.append((os.path.split(x["name"])[-1], int(x.get("size") or 0)))
if any(name != "consolidated.safetensors" for name, _ in shard_names):
shard_names = [x for x in shard_names if x[0] != "consolidated.safetensors"]
if not shard_names:
return
try:
for filename, _ in shard_names:
hf_hub_download(
repo_id = model_name,
filename = filename,
cache_dir = hub_cache_dir,
local_files_only = True,
token = token,
)
return # already fully cached
except Exception:
pass
# Mirror the merge's index filter (download path only): some repos ship leftover shards
# the index omits; keep only indexed ones, else the disk gate over-counts and we fetch
# unused shards.
if len(shard_names) > 1:
try:
import json as _json
_idx = hf_hub_download(
repo_id = model_name,
filename = "model.safetensors.index.json",
cache_dir = hub_cache_dir,
token = token,
)
with open(_idx, encoding = "utf-8") as _f:
_indexed = {
os.path.split(v)[-1] for v in _json.load(_f).get("weight_map", {}).values()
}
if _indexed and not {n for n, _ in shard_names}.issubset(_indexed):
_kept = [x for x in shard_names if x[0] in _indexed]
if _kept:
shard_names = _kept
except Exception:
pass
total_size_in_bytes = sum(size for _, size in shard_names)
# The cache copy is extra disk on top of the merge working copy; need room for both.
from huggingface_hub import constants as _hf_constants
# abspath so a relative HF_HUB_CACHE walks up to an existing root, not "".
cache_probe = os.path.abspath(
os.path.expanduser(str(hub_cache_dir or _hf_constants.HF_HUB_CACHE))
)
while cache_probe and not os.path.exists(cache_probe):
parent = os.path.dirname(cache_probe)
if parent == cache_probe:
break
cache_probe = parent
free_space = shutil.disk_usage(cache_probe).free if os.path.exists(cache_probe) else 0
if free_space < 2 * total_size_in_bytes:
print(
f"Unsloth: Not enough free disk to keep `{model_name}` in the Hugging Face "
f"cache (need ~{round(2 * total_size_in_bytes / 1024**3, 1)}GB free, have "
f"{round(free_space / 1024**3, 1)}GB). Downloading straight to the merge "
f"directory instead; the next export will re-download it."
)
return
if total_size_in_bytes >= 0.1 * 1024**3:
size_str = f"{round(total_size_in_bytes / 1024**3, 1)}GB"
else:
size_str = f"{max(1, round(total_size_in_bytes / 1024**2))}MB"
print(
f"Unsloth: Downloading `{model_name}` into the Hugging Face cache so future "
f"exports skip the {size_str} download..."
)
snapshot_download(
repo_id = model_name,
allow_patterns = [name for name, _ in shard_names]
+ ["model.safetensors.index.json", "tokenizer.model"],
cache_dir = hub_cache_dir,
token = token,
)
except Exception as e:
print(
f"Unsloth: Could not pre-cache the base model weights ({e}). "
f"Falling back to downloading into the merge directory."
)
@torch.inference_mode
def save_to_gguf_generic(
model,
@ -3888,6 +4064,7 @@ def unsloth_generic_save(
print(f"Unsloth: Model saved successfully to '{save_directory}'")
else:
_prewarm_base_model_hub_cache(model, save_method = save_method, token = token)
merge_and_overwrite_lora(
get_model_name,
model = model,

102
unsloth_cli/commands/start.py
View file

@ -983,7 +983,9 @@ def _session_config(agent: str, launch: bool):
else:
# Never wipe this dir: a previously printed recipe may still be running
# an agent whose sessions/state live here, and every config writer
# merges idempotently into an existing home anyway.
# merges idempotently into an existing home anyway. Writers must also
# reset any state a previous run's flags left behind (--yolo especially),
# since files here outlive the invocation that wrote them.
path = _agents_config_root() / agent
path.mkdir(parents = True, exist_ok = True, mode = 0o700)
yield path
@ -1043,6 +1045,62 @@ def write_openclaw_config(
{"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}},
)
typer.echo(f"Updated {approvals}")
else:
# The no-launch config dir is reused across runs, so a previous --yolo run may
# have left auto-approval state behind. OpenClaw treats an omitted exec policy as
# security=full, ask=off on the gateway host, so deleting the keys would keep
# auto-approval on: a non-yolo run must WRITE a prompting policy. Only a
# permissive/yolo policy is replaced; a stricter one set by hand survives.
tools = config.get("tools")
exec_policy = tools.get("exec") if isinstance(tools, dict) else None
exec_policy = exec_policy if isinstance(exec_policy, dict) else {}
# Match ONLY the exact fingerprint --yolo writes (host=gateway, security=full,
# ask=off, all explicit, no mode); anything else is left untouched. host=auto or an
# omitted host resolves to security=deny under an active sandbox, so treating those
# as the permissive gateway default would broaden a fresh sandboxed config from
# deny to allowlist. host=node and host=sandbox are user-set (--yolo only writes
# gateway). tools.exec.mode is OpenClaw's normalized knob (it cannot be combined
# with security/ask, and OpenClaw never rewrites our security/ask write into it),
# so a mode is always a deliberate user policy; never clobber it.
permissive = (
"mode" not in exec_policy
and exec_policy.get("host") == "gateway"
and exec_policy.get("security") == "full"
and exec_policy.get("ask") == "off"
)
if permissive:
exec_policy = _subdict(_subdict(config, "tools"), "exec")
exec_policy.pop("host", None) # routing only; defaults to the gateway host
exec_policy["security"] = "allowlist" # only allowlisted commands skip approval
exec_policy["ask"] = "on-miss" # prompt on every non-allowlisted command
# Drop the yolo defaults from the host approvals file (a stricter default set by
# the user or OpenClaw is kept). With a prompting tools.exec the stricter of the
# two layers wins, so an omitted approvals default still prompts.
approvals = path.parent / "exec-approvals.json"
if approvals.exists():
state = _read_json_object(approvals)
if state is not None:
defaults = state.get("defaults")
# Strip the defaults only when they are exactly the yolo fingerprint; a
# user-managed mixed policy that merely shares a field (e.g. askFallback=full,
# whose omitted default is deny) must be kept intact.
yolo_defaults = (("security", "full"), ("ask", "off"), ("askFallback", "full"))
is_yolo = isinstance(defaults, dict) and all(
defaults.get(k) == v for k, v in yolo_defaults
)
if is_yolo:
for k, _ in yolo_defaults:
del defaults[k]
if not defaults:
del state["defaults"]
if set(state) <= {"version"}:
# Nothing left but our own yolo payload: remove it.
approvals.unlink()
typer.echo(f"Removed {approvals}")
else:
# Keep approvals OpenClaw itself recorded; only the yolo defaults go.
_write_private_json(approvals, state)
typer.echo(f"Updated {approvals}")
if json.dumps(config, sort_keys = True) != before:
_write_private_json(path, config)
typer.echo(f"Updated {path}")
@ -1054,7 +1112,7 @@ def write_opencode_config(
model: dict,
path: Path,
yolo: bool = False,
) -> None:
) -> dict:
config = _read_json_object(path)
if config is None:
typer.echo(
@ -1062,7 +1120,7 @@ def write_opencode_config(
"yourself, or move the file aside and re-run.",
err = True,
)
return
return {}
before = json.dumps(config, sort_keys = True)
config.setdefault("$schema", "https://opencode.ai/config.json")
model_entry = {"name": model["id"]}
@ -1087,13 +1145,32 @@ def write_opencode_config(
compaction = _subdict(config, "compaction")
compaction["auto"] = True
compaction["reserved"] = max(1, window // 10)
tools = ("edit", "bash", "webfetch")
if yolo:
# OpenCode has no --yolo flag; auto-approve is the config `permission` block
# (singular). Allow the prompting tools so tool calls don't block on the TUI.
config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"}
# (singular). Allow the prompting tools so tool calls don't block on the TUI. This
# rides inline (OPENCODE_CONFIG_CONTENT) so --yolo works even over a project config.
session_permission = {t: "allow" for t in tools}
config["permission"] = dict(session_permission)
else:
# Undo only what --yolo wrote: our yolo sets an explicit per-tool "allow" for these
# three tools, so flip exactly those explicit allows back to "ask". A "deny"/"ask",
# a granular object, a string, or a "*" catch-all is the user's own rule and is left
# untouched. We do NOT carry a permission inline for a non-yolo session: since
# OPENCODE_CONFIG_CONTENT outranks the project opencode.json we cannot read, any
# value forced there would override the user's project rules (weakening a project
# deny, or auto-approving through a granular object's permissive default). Clearing
# our own persisted yolo state is the fix; the project's own permissions are honored.
session_permission: dict = {}
permission = config.get("permission")
if isinstance(permission, dict):
for tool in tools:
if permission.get(tool) == "allow":
permission[tool] = "ask"
if json.dumps(config, sort_keys = True) != before:
_write_private_json(path, config)
typer.echo(f"Updated {path}")
return session_permission
def write_hermes_config(base: str, model: dict, path: Path) -> None:
@ -1379,14 +1456,15 @@ def opencode(
# OPENCODE_CONFIG is an overlay (loaded between the user's global and project
# configs), so this adds the Unsloth provider/model for the session without
# changing the user's default model. Key lives in the config, not the env.
write_opencode_config(base, key, entry, config_path, yolo = yolo)
# A project's own opencode.json outranks OPENCODE_CONFIG, so the session model
# pin (and --yolo permissions) would silently lose to a repo config. Carry the
# settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project
# config; the API key stays in the private file, never in the printed env.
session_permission = write_opencode_config(base, key, entry, config_path, yolo = yolo)
# A project's own opencode.json outranks OPENCODE_CONFIG, so the session model pin
# would silently lose to a repo config. Carry it in OPENCODE_CONFIG_CONTENT, which
# outranks project config; the API key stays in the private file, never the env.
# Only --yolo carries a permission here (its allow must win over a project config);
# a non-yolo session returns no permission, so the project's own rules are honored.
inline_config: dict = {"model": f"unsloth/{entry['id']}"}
if yolo:
inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"}
if session_permission:
inline_config["permission"] = session_permission
env = {
"OPENCODE_CONFIG": str(config_path),
"OPENCODE_CONFIG_CONTENT": json.dumps(inline_config),

293
unsloth_cli/tests/test_start.py
View file

@ -446,7 +446,10 @@ def test_opencode_inline_config_beats_project_config(fake_studio):
assert "sk-unsloth" not in content_line # key stays in the private file
def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio):
def test_opencode_inline_config_omits_permission_without_yolo(fake_studio):
# A non-yolo session carries no permission inline. OPENCODE_CONFIG_CONTENT outranks the
# project opencode.json we cannot read, so forcing any value there would override the
# user's project rules; clearing our own config is the fix, and the inline pins the model.
result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert result.exit_code == 0, result.output
content_line = next(
@ -455,7 +458,8 @@ def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio):
inline = json.loads(
shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0]
)
assert inline == {"model": f"unsloth/{MODEL['id']}"}
assert inline["model"] == f"unsloth/{MODEL['id']}"
assert "permission" not in inline
def test_https_loopback_never_auto_serves(fake_studio, monkeypatch):
@ -1676,9 +1680,31 @@ def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path):
result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert result.exit_code == 0, result.output
config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text())
# A non-yolo run on a fresh config writes no permission block; it only flips a prior
# --yolo run's explicit allow back to ask (see the yolo-then-plain test below).
assert "permission" not in config
def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path):
# The core reset: a --yolo run wrote explicit per-tool allow; a later non-yolo run
# must flip exactly those back to ask so nothing stays auto-approved.
yolo = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"])
assert yolo.exit_code == 0, yolo.output
config_path = tmp_path / "agents" / "opencode" / "opencode.json"
assert json.loads(config_path.read_text())["permission"] == {
"edit": "allow",
"bash": "allow",
"webfetch": "allow",
}
plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert plain.exit_code == 0, plain.output
assert json.loads(config_path.read_text())["permission"] == {
"edit": "ask",
"bash": "ask",
"webfetch": "ask",
}
def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path):
result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"])
assert result.exit_code == 0, result.output
@ -1691,12 +1717,17 @@ def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path):
assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"}
def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path):
def test_no_yolo_openclaw_leaves_fresh_config_untouched(fake_studio, tmp_path):
# A fresh non-yolo run only undoes state a prior --yolo wrote; with no yolo
# fingerprint present it must not synthesize an exec policy. An omitted policy can
# resolve to a sandbox default of security=deny, so writing allowlist here would
# BROADEN it. The reset is scoped to the exact yolo write, verified by the
# yolo-then-plain round trip below.
result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"])
assert result.exit_code == 0, result.output
state = tmp_path / "agents" / "openclaw"
config = json.loads((state / "openclaw.json").read_text())
assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo
assert "exec" not in config.get("tools", {})
assert not (state / "exec-approvals.json").exists()
@ -1719,6 +1750,260 @@ def test_write_openclaw_config_yolo_unit(tmp_path):
}
def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp_path):
# The no-launch config dir is reused across runs, so a --yolo run persists its
# auto-approve settings; a later run without --yolo must strip them, not leave
# tool execution silently pre-approved.
yolo = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"])
assert yolo.exit_code == 0, yolo.output
config_path = tmp_path / "agents" / "opencode" / "opencode.json"
assert "permission" in json.loads(config_path.read_text())
plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert plain.exit_code == 0, plain.output
config = json.loads(config_path.read_text())
# The yolo allow policy is replaced by a prompting one, not deleted (which would
# revert to OpenCode's permissive "allow" default).
assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"}
# The session provider survives the cleanup.
assert "unsloth" in config["provider"]
def test_no_launch_rerun_clears_stale_openclaw_yolo_state(fake_studio, tmp_path):
yolo = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"])
assert yolo.exit_code == 0, yolo.output
state = tmp_path / "agents" / "openclaw"
assert (state / "exec-approvals.json").exists()
plain = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"])
assert plain.exit_code == 0, plain.output
config = json.loads((state / "openclaw.json").read_text())
# The yolo policy is replaced by a prompting one, not deleted (which would revert
# to OpenClaw's permissive default), and the yolo approvals file is gone.
assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"}
assert not (state / "exec-approvals.json").exists()
# The session provider survives the cleanup.
assert "unsloth" in config["models"]["providers"]
def test_write_openclaw_config_yolo_then_plain_unit(tmp_path):
path = tmp_path / "openclaw.json"
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True)
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
# A plain rerun replaces the yolo policy with a prompting one (deleting it would
# fall back to OpenClaw's permissive default) and removes the yolo approvals file.
assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"}
assert not (path.parent / "exec-approvals.json").exists()
def test_write_opencode_config_yolo_then_plain_unit(tmp_path):
path = tmp_path / "opencode.json"
start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True)
start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
# A plain rerun replaces the yolo allow policy with a prompting one.
assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"}
def test_openclaw_non_yolo_keeps_runtime_approvals(tmp_path):
# OpenClaw records its own entries in exec-approvals.json (OPENCLAW_STATE_DIR is
# this dir); the non-yolo reset drops only the yolo defaults, not those.
path = tmp_path / "openclaw.json"
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True)
approvals = path.parent / "exec-approvals.json"
state = json.loads(approvals.read_text())
state["agents"] = {"main": {"allowlist": ["git status"]}}
approvals.write_text(json.dumps(state))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
remaining = json.loads(approvals.read_text())
assert "defaults" not in remaining
assert remaining["agents"] == {"main": {"allowlist": ["git status"]}}
def test_openclaw_non_yolo_keeps_mixed_approval_defaults(tmp_path):
# A mixed user-managed defaults block that only shares a field with the yolo payload
# (here askFallback=full, whose omitted default is deny) is not stale yolo state, so a
# non-yolo run leaves it intact rather than stripping the shared field.
path = tmp_path / "openclaw.json"
approvals = path.parent / "exec-approvals.json"
mixed = {
"version": 1,
"defaults": {"security": "allowlist", "ask": "on-miss", "askFallback": "full"},
}
approvals.write_text(json.dumps(mixed))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
assert json.loads(approvals.read_text()) == mixed
def test_openclaw_non_yolo_leaves_partial_policy_untouched(tmp_path):
# A policy that lacks the full yolo fingerprint (here no host and no security) is not
# our --yolo write, so a non-yolo run leaves it as-is rather than assuming ask=off
# means permissive: an omitted host/security can resolve to a sandbox deny default.
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": {"timeout": 30, "ask": "off"}}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"timeout": 30, "ask": "off"}
def test_openclaw_non_yolo_leaves_no_permissive_values(tmp_path):
# The whole point of the reset: after a yolo run, a plain run must leave neither the
# config nor the approvals file at OpenClaw's permissive (security=full, ask=off)
# default, or exec still auto-approves.
path = tmp_path / "openclaw.json"
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True)
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
exec_policy = json.loads(path.read_text())["tools"]["exec"]
assert exec_policy.get("security") != "full"
assert exec_policy.get("ask") != "off"
assert not (path.parent / "exec-approvals.json").exists()
def test_openclaw_non_yolo_preserves_stricter_exec_policy(tmp_path):
# A policy that doesn't carry the yolo values (for example stricter security or
# prompting turned on) was not written by --yolo and must survive a plain run.
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": {"security": "deny", "ask": "on"}}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"security": "deny", "ask": "on"}
def test_openclaw_non_yolo_preserves_stricter_approval_defaults(tmp_path):
# exec-approvals.json defaults that don't match the yolo payload (stricter
# settings from the user or the OpenClaw UI) are kept, and the file stays.
path = tmp_path / "openclaw.json"
approvals = path.parent / "exec-approvals.json"
approvals.write_text(
json.dumps({"version": 1, "defaults": {"security": "allowlist", "ask": "on"}})
)
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
state = json.loads(approvals.read_text())
assert state["defaults"] == {"security": "allowlist", "ask": "on"}
def test_openclaw_non_yolo_leaves_unparseable_approvals(tmp_path):
# An unreadable approvals file is left in place rather than deleted, matching
# how an unparseable config is handled.
path = tmp_path / "openclaw.json"
approvals = path.parent / "exec-approvals.json"
approvals.write_text("{not json")
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
assert approvals.read_text() == "{not json"
def test_opencode_non_yolo_flips_only_explicit_allow(tmp_path):
# Only a tool explicitly set to "allow" (what --yolo writes) is flipped to "ask". A
# deny/ask a user set is kept, and an absent tool is not added.
path = tmp_path / "opencode.json"
path.write_text(json.dumps({"permission": {"edit": "allow", "bash": "deny", "read": "ask"}}))
session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["permission"] == {"edit": "ask", "bash": "deny", "read": "ask"}
assert session == {} # a non-yolo session carries no permission inline
def test_opencode_non_yolo_leaves_string_permission(tmp_path):
# A global string rule ("deny") is a user-managed catch-all; leave it untouched and
# carry no inline override.
path = tmp_path / "opencode.json"
path.write_text(json.dumps({"permission": "deny"}))
session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
assert json.loads(path.read_text())["permission"] == "deny"
assert session == {}
def test_opencode_non_yolo_leaves_catch_all_and_flips_explicit_allow(tmp_path):
# A "*" catch-all is the user's own rule, never something --yolo writes (yolo sets
# explicit per-tool allow), so it is left intact; an explicit per-tool "allow" is still
# flipped to "ask", but an absent tool inheriting the catch-all is not touched.
path = tmp_path / "opencode.json"
path.write_text(json.dumps({"permission": {"*": "allow", "bash": "allow"}}))
session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
assert json.loads(path.read_text())["permission"] == {"*": "allow", "bash": "ask"}
assert session == {}
def test_opencode_non_yolo_leaves_granular_object(tmp_path):
# A granular object value is a user rule (yolo only ever writes a plain "allow" string),
# so it is left in the file verbatim and never carried inline.
path = tmp_path / "opencode.json"
obj = {"read *": "deny", "git *": "ask"}
path.write_text(json.dumps({"permission": {"bash": dict(obj)}}))
session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
assert json.loads(path.read_text())["permission"]["bash"] == obj
assert session == {}
def test_openclaw_non_yolo_leaves_mode_policy(tmp_path):
# tools.exec.mode is OpenClaw's normalized knob and cannot be combined with explicit
# security/ask (the config is rejected), so a mode-based policy must be left as-is.
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": {"mode": "deny"}}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"mode": "deny"}
def test_openclaw_non_yolo_preserves_sandbox_host(tmp_path):
# host=sandbox defaults to security=deny (stricter than the gateway "full" default),
# so a non-yolo run must not treat the missing security as permissive nor pop host
# (which would broaden routing to the gateway).
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": {"host": "sandbox"}}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"host": "sandbox"}
def test_openclaw_non_yolo_preserves_node_host(tmp_path):
# host=node routes to a paired node and is only ever set by the user (--yolo writes
# host=gateway), so a non-yolo run must not pop it and reroute to the gateway.
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": {"host": "node"}}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"host": "node"}
def test_openclaw_non_yolo_preserves_auto_host_permissive(tmp_path):
# host=auto (or omitted) with security=full/ask=off is NOT the --yolo write: under an
# active sandbox, auto resolves to security=deny. --yolo only ever writes host=gateway,
# so the reset must not treat auto/None as the permissive gateway default and broaden a
# sandboxed deny to allowlist.
for exec_policy in (
{"host": "auto", "security": "full", "ask": "off"},
{"security": "full", "ask": "off"},
):
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": dict(exec_policy)}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == exec_policy
def test_openclaw_non_yolo_resets_only_gateway_yolo_fingerprint(tmp_path):
# The reset fires on exactly the host=gateway + security=full + ask=off write --yolo
# makes, and nothing else.
path = tmp_path / "openclaw.json"
path.write_text(
json.dumps({"tools": {"exec": {"host": "gateway", "security": "full", "ask": "off"}}})
)
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"}
def test_openclaw_non_yolo_preserves_full_mode(tmp_path):
# OpenClaw never normalizes our security=full/ask=off yolo write into mode:"full"
# (verified against the binary: doctor --fix and config get leave security/ask as-is),
# so a mode:"full" is always a deliberate user policy, not stale yolo state; leave it.
path = tmp_path / "openclaw.json"
path.write_text(json.dumps({"tools": {"exec": {"mode": "full"}}}))
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
assert config["tools"]["exec"] == {"mode": "full"}
def test_yolo_command_flags_unmapped_agent_is_empty():
# Config-based agents (and any typo) must yield no flag, not a KeyError.
assert start._yolo_command_flags("opencode", True) == []