diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 9b85b20177..d430d2c172 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -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. diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000000..f9a270540f --- /dev/null +++ b/.github/workflows/ossf.yml @@ -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 diff --git a/install.ps1 b/install.ps1 index 8c667df079..9114f80af9 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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) diff --git a/install.sh b/install.sh index 81c50bc899..796d80e401 100755 --- a/install.sh +++ b/install.sh @@ -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..." diff --git a/pyproject.toml b/pyproject.toml index 844ead2454..80b3d757e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh index 5ef9ee386a..aa560fc432 100644 --- a/scripts/install_rocm_wsl_strixhalo.sh +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -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 <>> 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() diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1c7a409bc1..0633f80bbc 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -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", diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 7b572a28ff..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -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 minus 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 ```` delta when the streamed text left a ```` + 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": ""}, + }, + ) + ] + 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("") - new_text.count("") 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]: diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index b64605e16f..a1d03c03e0 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -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", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5e6287f528..757467a008 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 block) then the held + """Close a live-streamed 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 += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed 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 += "" + 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 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 += "" 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 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 ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 19d2230278..cf5d24c367 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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: diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 345b4dd853..47d26209b4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -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) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 963d584303..7a27a58a52 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -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: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 862bce8be8..914699f540 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -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 diff --git a/studio/backend/run.py b/studio/backend/run.py index ccd0113972..2cc6c4a93e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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 diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 41a9adcc29..87aa50ee26 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -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: diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 170b456eac..0c6550a3bb 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -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 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 . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "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) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index aa19df15fe..0239410734 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -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()) diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 5d14c5c5bd..8408f8203d 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -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) diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py new file mode 100644 index 0000000000..19808ad0d7 --- /dev/null +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -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 -%}{{- '' -}}{%- else -%}{{- '' -}}{%- 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 = "", 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 , 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("") + assert "Absolute maximum" not in out + + +def test_high_state_renders_plain_thinking(): + """UI 'High' -> et=true, effort=high -> open , 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("") + assert "Absolute maximum" not in out + + +def test_max_state_injects_max_preamble(): + """UI 'Max' -> et=true, effort=max -> open 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("") + 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("") + assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py new file mode 100644 index 0000000000..940b35d7ba --- /dev/null +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -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"]) diff --git a/studio/backend/tests/test_inference_default_models_non_blocking.py b/studio/backend/tests/test_inference_default_models_non_blocking.py new file mode 100644 index 0000000000..83a8e7bbfb --- /dev/null +++ b/studio/backend/tests/test_inference_default_models_non_blocking.py @@ -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 diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fb1b0e52b7..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -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 . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.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"] - == "I am thinking. Still thinking.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 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("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.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] == "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 + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . 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] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +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 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 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] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # 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 = [ diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index ccbd78e2b1..05e017ba7f 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -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 diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 9fd1535f22..0ed670ac01 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -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 diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index db66df8a30..b5f1069f12 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -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) diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 463d26a692..9d2c983aed 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -487,7 +487,7 @@ TEMPLATE_TO_RESPONSES_MAPPER = { }, "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index c2e389ca5f..1aff5793b2 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -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 { + 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() { Export - {[ - { label: "Raw JSONL", fn: exportConversationRawJsonl }, - { label: "CSV", fn: exportConversationCsv }, - { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, - ].map(({ label, fn }) => ( + {CHAT_EXPORT_OPTIONS.map(({ label, format }) => ( { @@ -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"), )} + {/* "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 && ( +

+ {t("shell.navigation.noChatsYet")} +

+ )} diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx new file mode 100644 index 0000000000..823696693a --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -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; +}; + +type ContextUsageMetadata = { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + cachedTokens?: number; + cacheWriteTokens?: number; + modelId?: string; +}; + +type MessageCustomMetadata = { + responseDetails?: ResponseDetailsMetadata; + contextUsage?: ContextUsageMetadata; + serverTimings?: Record; + 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 = { + search: "Search", + fetch: "Fetch", + code: "Code", + images: "Images", + mcp: "MCP", + docs: "Docs", + artifacts: "Canvas", +}; + +const TOOL_CALL_LABELS: Record = { + 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 | undefined, + toolCalls: string[], +): string | null { + if (!tools && toolCalls.length === 0) return null; + const activeKeys = new Set(); + 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 ( +
+

{title}

+
{children}
+
+ ); +} + +function DetailRow({ + label, + value, + mono = false, +}: { + label: string; + value: ReactNode | null | undefined; + mono?: boolean; +}) { + if (value == null || value === "") return null; + return ( +
+ {label} + + {value} + +
+ ); +} + +function useResponseModelDisplay() { + const message = useMessage(); + const models = useChatRuntimeStore((s) => s.models); + const providers = useExternalProvidersStore((s) => s.providers); + + const custom = ( + message.metadata as Record | 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 ( + + {modelLabel} + + ); +}; + +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 ( + + + + + + Response details + + + Timing, model, token, and tool details for this response. + + + +
+
+

+ {summaryLabel} +

+ {providerLabel ? ( +

+ {providerLabel} +

+ ) : null} +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+ + + ); +}; diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index c5b53577cb..96d21d6fe7 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -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} > -
+
-
+ + + +
{isOpen && !isReasoningStreaming && ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 7fe98b701f..09551cd413 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -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 (
@@ -3649,6 +3656,11 @@ const AssistantMessage: FC = () => {
) : ( <> + {!hasReasoningParts ? ( +
+ +
+ ) : null} @@ -3893,58 +3905,76 @@ const EditAssistantMessageButton: FC = () => { const AssistantActionBar: FC = () => { const { forkMessage, forkDisabled } = useForkMessageAction(); + const [detailsOpen, setDetailsOpen] = useState(false); return ( - - - - - - - - - - - - - - + <> + + + + + + - - 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" - > - 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" + + + + + + + + + + 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" > - - Fork in new chat - - - + 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" + > + + Fork in new chat + + + + + Export as Markdown + + + 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" + > - Export as Markdown + See response details - - - - - + + + + + + ); }; diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index def1b9bad9..63cdd03141 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -54,6 +54,17 @@ export const usePlatformStore = create()((_, 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 }); diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index df266fd749..2ca6ceddf7 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -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[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, }, }, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index cd7cfc77fc..b155eff780 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -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({
{projectTab === "sources" ? ( - + + Loading sources… +
+ } + > + + ) : (
{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). diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e23d1b0b33..798c1658f9 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -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); diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index 55d7777e43..0a0df1139b 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -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(); + const pairItems = new Map(); 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([]); + 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 { diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 75749b0040..7cd9611c71 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -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"; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 67980b94c6..b545695f9e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -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 { - 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 { - 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, diff --git a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts index f019b7dc3d..5b5011a78f 100644 --- a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts @@ -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()( @@ -23,6 +26,9 @@ export const useChatPreferencesStore = create()( showModelDisclaimer: true, setShowModelDisclaimer: (showModelDisclaimer) => set({ showModelDisclaimer }), + showResponseModel: false, + setShowResponseModel: (showResponseModel) => + set({ showResponseModel }), }), { name: "unsloth_chat_preferences", @@ -32,6 +38,7 @@ export const useChatPreferencesStore = create()( ...current, confirmDeleteChats: saved?.confirmDeleteChats ?? true, showModelDisclaimer: saved?.showModelDisclaimer ?? true, + showResponseModel: saved?.showResponseModel ?? false, }; }, }, diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 3fe69ccc26..111510d925 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -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 diff --git a/studio/frontend/src/features/chat/utils/chat-history-storage.ts b/studio/frontend/src/features/chat/utils/chat-history-storage.ts index 4294887df0..00df3657b1 100644 --- a/studio/frontend/src/features/chat/utils/chat-history-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-history-storage.ts @@ -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( diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts index 8b6bc7ee7f..9a61142f73 100644 --- a/studio/frontend/src/features/settings/api/embedding-model.ts +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -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"), diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx index b26b09e022..9519898b21 100644 --- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx @@ -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} /> + + + diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 7670aae5fa..8fd70cc1b7 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -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( diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index f428567cee..93ae9ad70d 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -43,6 +43,7 @@ export const en = { video: "Video", export: "Export", recents: "Recents", + noChatsYet: "No chats yet", settings: "Settings", api: "API", lightMode: "Light Mode", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index c9be6cce5f..38b74bb3d4 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -43,6 +43,7 @@ export const zhCN = { video: "视频", export: "导出", recents: "最近", + noChatsYet: "暂无对话", settings: "设置", api: "API", lightMode: "浅色模式", diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 1a4ebdace8..86a9634048 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -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 = /(? { 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. `(? 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; diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx index 0922e764bb..d0ddf2fc6e 100644 --- a/studio/frontend/src/main.tsx +++ b/studio/frontend/src/main.tsx @@ -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( - - - , - ); -}); +createRoot(rootElement).render( + + + , +); + +fetchDeviceType().catch(() => undefined); diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 5a48d26632..7ef5244754 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -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> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + + #[cfg(unix)] + struct ManagedCapabilityCacheHome { + path: PathBuf, + previous: Option, + } + + #[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"; diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 57b8365ec5..0d20f271c5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -188,6 +188,16 @@ fn managed_bin_fingerprint(bin: &Path) -> Option { } fn capability_cache_path() -> Option { + #[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", diff --git a/tests/_zoo_rocm_spoof.py b/tests/_zoo_rocm_spoof.py new file mode 100644 index 0000000000..050191e9d1 --- /dev/null +++ b/tests/_zoo_rocm_spoof.py @@ -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) diff --git a/tests/saving/test_prewarm_base_model_hub_cache.py b/tests/saving/test_prewarm_base_model_hub_cache.py new file mode 100644 index 0000000000..cbb52863ba --- /dev/null +++ b/tests/saving/test_prewarm_base_model_hub_cache.py @@ -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 diff --git a/tests/studio/install/test_rocm_rdna_routing.py b/tests/studio/install/test_rocm_rdna_routing.py new file mode 100644 index 0000000000..b4aeafb7e4 --- /dev/null +++ b/tests/studio/install/test_rocm_rdna_routing.py @@ -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] diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index c8f2053946..bfc8132683 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -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.""" diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py new file mode 100644 index 0000000000..04301a0de6 --- /dev/null +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -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 diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 98f7e2db62..7a738e236c 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 169b610988..1c75f8ce66 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.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}( diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c25a031b82..a1da099758 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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 diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 62ef9e916a..b5cadf2dea 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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" diff --git a/unsloth/save.py b/unsloth/save.py index 020c63a9e2..0b408a1889 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -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, diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index b188180188..1895125b11 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -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), diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index a6a092a17c..affc24626e 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -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) == []