diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000000..17d96cd0f5 --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,8 @@ +# Commits listed here are skipped by `git blame` so that bulk, whitespace-only +# changes don't obscure the real authorship of a line. +# +# GitHub honors this file automatically. To use it locally, run once: +# git config blame.ignoreRevsFile .git-blame-ignore-revs + +# chore(studio/frontend): normalize line endings to LF +c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a diff --git a/.gitattributes b/.gitattributes index 75fba5d6ab..5f04b5e9d1 100644 --- a/.gitattributes +++ b/.gitattributes @@ -5,3 +5,9 @@ # clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks # them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). *.sh text eol=lf + +# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather +# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files +# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts) +# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF. +studio/frontend/** text=auto eol=lf diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml new file mode 100644 index 0000000000..4632794587 --- /dev/null +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs tests/python/test_cross_platform_parity.py on Windows and macOS. +# +# Why: that test is the guard that install.sh and install.ps1 stay in +# sync, but today it only runs on ubuntu-latest (auto-discovered by +# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both +# installer scripts, and on Windows Path.read_text() defaults to the +# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already +# contains a U+274C) raises UnicodeDecodeError there even though Linux and +# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this job keeps that from silently regressing by exercising the +# test on the platforms it claims parity for. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. + +name: Cross-platform parity + +on: + pull_request: + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + push: + branches: [main] + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + parity: + name: parity (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 10 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - run: python -m pip install -U pip pytest + - name: Cross-platform parity test + run: python -m pytest tests/python/test_cross_platform_parity.py -q diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index eb60a5a201..6eb8d1bc6e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,3 +27,9 @@ Your support extends beyond code: Finally, please be mindful of our [Code of Conduct](https://github.com/unslothai/unsloth/blob/main/CODE_OF_CONDUCT.md) to ensure a welcoming and inclusive environment for everyone. Thank you so much for reading and we hope you have lots of fun using Unsloth! 🦥 + + +## Pull Request Guidelines +- Keep PRs focused on a single change +- Include a concise description and motivation +- Link related issues when applicable diff --git a/install.ps1 b/install.ps1 index fd7d3f4a81..9abddc9ce2 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1177,14 +1177,38 @@ shell.Run cmd, 0, False if ($SkipTorch) { $InitialGpuBranch = "no_torch" } Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion - # ── Install uv if not present ── + # ── Install uv ── Write-TauriLog "STEP" "Installing uv package manager" - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - substep "installing uv package manager..." + $UvMinVersion = "0.7.22" + function Test-UvVersionOk { + $cmd = Get-Command uv -ErrorAction SilentlyContinue + if (-not $cmd) { return $false } + try { + $raw = (& uv --version 2>$null | Select-Object -First 1) + } catch { + return $false + } + if ($raw -notmatch 'uv\s+([0-9]+(?:\.[0-9]+)+)') { return $false } + try { + return ([version]$Matches[1] -ge [version]$UvMinVersion) + } catch { + return $false + } + } + + if (-not (Test-UvVersionOk)) { + if (Get-Command uv -ErrorAction SilentlyContinue) { + substep "updating uv package manager..." + } else { + substep "installing uv package manager..." + } if ($script:WingetAvailable) { $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" - try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + try { winget upgrade --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + if (-not (Test-UvVersionOk)) { + try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + } $ErrorActionPreference = $prevEAP Refresh-SessionPath } @@ -1192,19 +1216,40 @@ shell.Run cmd, 0, False # use Astral's official PowerShell installer. This is the only # supported path on hosts without winget (Windows ARM64 runners, # corporate machines without the Store, etc.). - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { + if (-not (Test-UvVersionOk)) { substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow" Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") Refresh-SessionPath } } - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { + # A freshly installed uv can sit later on PATH than an older one (active + # venv, Scoop/pipx shim). Prefer a just-installed uv from a known location. + if (-not (Test-UvVersionOk)) { + $origPath = $env:PATH + foreach ($d in @($env:UV_INSTALL_DIR, $env:XDG_BIN_HOME, + (Join-Path $env:USERPROFILE ".local\bin"), + (Join-Path $env:LOCALAPPDATA "Microsoft\WinGet\Links"))) { + if ($d -and (Test-Path $d)) { + $env:PATH = "$d;$origPath" + if (Test-UvVersionOk) { break } + $env:PATH = $origPath + } + } + } + + if (-not (Test-UvVersionOk)) { step "uv" "could not be installed" "Red" substep "Install it from https://docs.astral.sh/uv/" "Yellow" return (Exit-InstallFailure "uv could not be installed") } + # When bytecode compilation is enabled, large installs can exceed uv's 60s + # default on slow machines. Default to 180s, preserving overrides ("0" disables). + if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT) { + $env:UV_COMPILE_BYTECODE_TIMEOUT = "180" + } + # ── Create venv (migrate old layout if possible, otherwise fresh) ── # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. @@ -1876,7 +1921,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1890,7 +1935,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1937,7 +1982,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1949,7 +1994,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1977,7 +2022,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.2" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.3" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index df9eca65c7..af062ae057 100755 --- a/install.sh +++ b/install.sh @@ -1456,7 +1456,11 @@ fi # ── Install uv ── tauri_log "STEP" "Installing uv package manager" -UV_MIN_VERSION="0.7.14" +UV_MIN_VERSION="0.7.22" + +# When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables). +: "${UV_COMPILE_BYTECODE_TIMEOUT:=180}" +export UV_COMPILE_BYTECODE_TIMEOUT version_ge() { # returns 0 if $1 >= $2 @@ -2043,6 +2047,37 @@ _pick_radeon_wheel() { # CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts # the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 + # librocdxg), then sources the env it persisted so detection finds the GPU. +# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d +# so non-login Studio/llama launches inherit it. Idempotent (writes only when +# the drop-in is missing); no-op without librocdxg, so never fires off WSL. +# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes +# after this shell on a non-root reinstall. Best-effort either way. +_persist_rocm_wsl_dropin() { + [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ] || return 0 + _rw_rocm=/opt/rocm + export HSA_ENABLE_DXG_DETECTION=1 + export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 + case ":${PATH}:" in + *":${_rw_rocm}/bin:"*) ;; + *) export PATH="${_rw_rocm}/bin:${PATH}" ;; + esac + export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" + [ -r /etc/profile.d/unsloth-rocm-wsl.sh ] && return 0 + _rw_dropin="$( + printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' + printf 'export HSA_ENABLE_DXG_DETECTION=1\n' + printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' + printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" + printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" + printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' + )" + if [ "$(id -u)" = "0" ]; then + printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true + elif command -v sudo >/dev/null 2>&1; then + printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true + fi +} + _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 @@ -2056,6 +2091,11 @@ _maybe_bootstrap_rocm_wsl() { _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + # rocminfo may work only via the transient env _ensure_rocm_probe_env + # just set, which dies with the installer. Persist the drop-in so login + # shells (Studio, llama.cpp) inherit it -- else a reinstall over an + # existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU. + _persist_rocm_wsl_dropin return 0 fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). @@ -2073,31 +2113,8 @@ _maybe_bootstrap_rocm_wsl() { . /etc/profile.d/unsloth-rocm-wsl.sh || true else # librocdxg present but the env drop-in is gone (e.g. a Studio - # uninstall removed it while keeping shared ROCm). Restore the FULL - # env inline (so rocminfo is on PATH) and recreate the drop-in. - _rw_rocm=/opt/rocm - export HSA_ENABLE_DXG_DETECTION=1 - export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 - export PATH="${_rw_rocm}/bin:${PATH}" - export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" - # Persist the drop-in so later non-login Studio launches get the env - # too. /etc/profile.d is root-owned: a plain redirect fails for a - # non-root reinstall (ROCm would silently disappear after this shell), - # so tee through sudo when not root. Best-effort -- the current shell - # already has the env, so the install proceeds either way. - _rw_dropin="$( - printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' - printf 'export HSA_ENABLE_DXG_DETECTION=1\n' - printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' - printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" - printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" - printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' - )" - if [ "$(id -u)" = "0" ]; then - printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true - elif command -v sudo >/dev/null 2>&1; then - printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true - fi + # uninstall removed it while keeping shared ROCm). Restore the env. + _persist_rocm_wsl_dropin fi return 0 fi @@ -2405,7 +2422,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.2" unsloth-zoo + "unsloth>=2026.6.3" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2418,7 +2435,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.2" unsloth-zoo + "unsloth>=2026.6.3" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2622,7 +2639,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.2" unsloth-zoo + "unsloth>=2026.6.3" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2640,7 +2657,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.2" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2672,7 +2689,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.2" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.3" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." diff --git a/studio/backend/assets/chat_templates/gemma-4-edge.jinja b/studio/backend/assets/chat_templates/gemma-4-edge.jinja new file mode 100644 index 0000000000..0266127233 --- /dev/null +++ b/studio/backend/assets/chat_templates/gemma-4-edge.jinja @@ -0,0 +1,397 @@ +{#- + Gemma 4 chat template (E2B / E4B edge variant), vendored for Unsloth Studio. + Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking + flag plus null-rendering, string-arguments validation, balanced turn tags, empty + messages handling, and OpenAI image_url/input_audio aliases). + Studio-local changes vs PR #118: + 1. preserve_thinking defaults to false (see SETUP block below). + 2. The empty "<|channel>thought\n" block on enable_thinking=false is + NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it, + google/gemma-4-E4B-it) that omits it; only the 12b/26B-A4B/31B family emits it. + This file matches the E2B/E4B behavior; gemma-4.jinja keeps the larger-model one. + Applied to unsloth/gemma-4-E2B-it-GGUF and unsloth/gemma-4-E4B-it-GGUF so the + embedded GGUF template does not need re-downloading. +-#} +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%} + {%- if thinking_text and thinking_gate and message.get('tool_calls') -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and not message.get('tool_calls') + and not ns_tr_out.flag + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {{- '\n' -}} + {%- elif not (ns_tr_out.flag and not has_content) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- endif -%} + {#- E2B/E4B do NOT emit an empty thought block when enable_thinking is false + (unlike the 12b/26B-A4B/31B family); see header. -#} +{%- endif -%} diff --git a/studio/backend/assets/chat_templates/gemma-4.jinja b/studio/backend/assets/chat_templates/gemma-4.jinja new file mode 100644 index 0000000000..65ab39df57 --- /dev/null +++ b/studio/backend/assets/chat_templates/gemma-4.jinja @@ -0,0 +1,397 @@ +{#- + Gemma 4 chat template, vendored for Unsloth Studio. + Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking + flag plus null-rendering, string-arguments validation, balanced turn tags, empty + messages handling, and OpenAI image_url/input_audio aliases). + Studio-local change: preserve_thinking defaults to false (see SETUP block below). + Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not + need re-downloading. Keep in sync with upstream if PR #118 changes. +-#} +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%} + {%- if thinking_text and thinking_gate and message.get('tool_calls') -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and not message.get('tool_calls') + and not ns_tr_out.flag + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {{- '\n' -}} + {%- elif not (ns_tr_out.flag and not has_content) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- endif -%} + + {%- if not enable_thinking -%} + {#- Suppress thinking - but not when awaiting tool responses -#} + {%- if ns.prev_message_type != 'tool_call' -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} + {%- endif -%} +{%- endif -%} diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d5a3de06df..b28b61f088 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -18,7 +18,12 @@ from utils.hardware import clear_gpu_cache from utils.models import is_vision_model, get_base_model_from_lora from utils.models.model_config import detect_audio_type -from utils.paths import ensure_dir, outputs_root, resolve_export_dir, resolve_output_dir +from utils.paths import ( + ensure_dir, + outputs_root, + resolve_export_write_dir, + resolve_output_dir, +) from core.inference import get_inference_backend # GPU-only imports — guarded for Apple Silicon where these aren't needed @@ -336,7 +341,7 @@ class ExportBackend: save_method = "merged_16bit" if save_directory: - save_directory = str(resolve_export_dir(save_directory)) + save_directory = str(resolve_export_write_dir(save_directory)) logger.info(f"Saving merged model locally to: {save_directory}") ensure_dir(Path(save_directory)) @@ -436,7 +441,7 @@ class ExportBackend: output_path: Optional[str] = None try: if save_directory: - save_directory = str(resolve_export_dir(save_directory)) + save_directory = str(resolve_export_write_dir(save_directory)) logger.info(f"Saving base model locally to: {save_directory}") ensure_dir(Path(save_directory)) @@ -563,6 +568,7 @@ class ExportBackend: return False, "No model loaded. Please select a checkpoint first.", None output_path: Optional[str] = None + model_tmp_to_cleanup: Optional[str] = None try: # unsloth expects lowercase quant method quant_method = quantization_method.lower() @@ -588,9 +594,8 @@ class ExportBackend: _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True if save_directory: - save_directory = str(resolve_export_dir(save_directory)) - # Absolute path so unsloth's relative-path internals resolve - # against the repo root cwd, not the export directory. + save_directory = str(resolve_export_write_dir(save_directory)) + # Keep unsloth relative-path internals anchored to the repo cwd. abs_save_dir = os.path.abspath(save_directory) logger.info(f"Saving GGUF model locally to: {abs_save_dir}") @@ -604,9 +609,15 @@ class ExportBackend: cwd = os.getcwd() pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - model_save_path = os.path.join(abs_save_dir, "model") + pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()} + + # Avoid clobbering an existing user-owned model/ directory. + import uuid + + _model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}") + model_tmp_to_cleanup = _model_tmp self.current_model.save_pretrained_gguf( - model_save_path, + _model_tmp, self.current_tokenizer, quantization_method = quant_method, ) @@ -618,10 +629,12 @@ class ExportBackend: shutil.move(src, dest) logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/") - # Flatten any .gguf from subdirs (e.g. model_gguf/) into abs_save_dir. + # Flatten GGUF files from subdirs created during this export. for sub in list(Path(abs_save_dir).iterdir()): if not sub.is_dir(): continue + if sub.name in pre_existing_subs: + continue for src in sub.glob("*.gguf"): dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) @@ -634,7 +647,7 @@ class ExportBackend: if self.current_checkpoint: ckpt = Path(self.current_checkpoint) gguf_dir = ckpt.parent / f"{ckpt.name}_gguf" - if gguf_dir.is_dir(): + if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve(): for src in gguf_dir.glob("*.gguf"): dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) @@ -683,6 +696,8 @@ class ExportBackend: ) except Exception as e: + if model_tmp_to_cleanup: + shutil.rmtree(model_tmp_to_cleanup, ignore_errors = True) logger.error(f"Error exporting GGUF model: {e}") import traceback @@ -712,7 +727,7 @@ class ExportBackend: output_path: Optional[str] = None try: if save_directory: - save_directory = str(resolve_export_dir(save_directory)) + save_directory = str(resolve_export_write_dir(save_directory)) logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) diff --git a/studio/backend/core/inference/chat_templates.py b/studio/backend/core/inference/chat_templates.py new file mode 100644 index 0000000000..58f63ff61b --- /dev/null +++ b/studio/backend/core/inference/chat_templates.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Bundled chat-template selection for GGUF inference. + +Some shipped GGUF quants embed an older chat template. Rather than re-cutting and +asking users to re-download every quant, Studio can override the embedded template +at llama-server launch time with a bundled, up-to-date Jinja template for known +model families. The override is wired through the existing ``chat_template_override`` +-> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``. + +Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118 +``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking" +toggle appears while staying disabled by default. +""" + +import re +from functools import lru_cache +from pathlib import Path +from typing import Optional + +# assets live at /assets/chat_templates/. This module is at +# /core/inference/chat_templates.py, so walk up three parents to +# (mirrors utils/inference/inference_config.py). +_ASSETS_DIR = Path(__file__).parent.parent.parent / "assets" / "chat_templates" + +# unsloth/gemma-4--GGUF (case-insensitive). The "-GGUF" suffix is retained +# on ModelConfig.identifier for HF GGUF repos, so this matches E2B / E4B / 31B / +# 26B-A4B and any future unsloth/gemma-4-*-GGUF, while excluding gemma-3, +# non-Unsloth, and non-GGUF identifiers (e.g. the bf16 "unsloth/gemma-4-E2B-it"). +_GEMMA4_GGUF_RE = re.compile(r"^unsloth/gemma-4-.+-gguf$", re.IGNORECASE) + +# Google ships two distinct gemma-4 chat templates: E2B/E4B omit the empty +# "<|channel>thought" block on enable_thinking=false, while the +# 12b/26B-A4B/31B family emits it. Route the two GGUF families to the matching +# bundled template so each keeps its model's intended behavior. +_GEMMA4_EDGE_GGUF_RE = re.compile(r"^unsloth/gemma-4-e[24]b-it-gguf$", re.IGNORECASE) + +_GEMMA4_TEMPLATE_FILE = "gemma-4.jinja" # 12b / 26B-A4B / 31B +_GEMMA4_EDGE_TEMPLATE_FILE = "gemma-4-edge.jinja" # E2B / E4B + + +def _canonical_repo_id(model_identifier: str) -> str: + """Mirror ``ModelConfig.from_identifier``: a bare HF shorthand with no owner + (e.g. ``gemma-4-E2B-it-GGUF``) defaults to the ``unsloth/`` org. The resolver + runs on the raw ``request.model_path`` (before that canonicalization), so apply + the same rule here, otherwise shorthand loads would skip the override. + """ + mid = model_identifier.strip() + if mid and "/" not in mid: + mid = f"unsloth/{mid}" + return mid + + +def is_unsloth_gemma4_gguf(model_identifier: Optional[str]) -> bool: + """True for canonical ``unsloth/gemma-4-*-GGUF`` repo identifiers (and the + owner-less shorthand that resolves to the same Unsloth repo).""" + if not model_identifier: + return False + return bool(_GEMMA4_GGUF_RE.match(_canonical_repo_id(model_identifier))) + + +def is_unsloth_gemma4_edge_gguf(model_identifier: Optional[str]) -> bool: + """True for the E2B / E4B GGUF repos, which use the edge-variant template.""" + if not model_identifier: + return False + return bool(_GEMMA4_EDGE_GGUF_RE.match(_canonical_repo_id(model_identifier))) + + +def _gemma4_template_file(model_identifier: Optional[str]) -> Optional[str]: + """Return the bundled template filename for a gemma-4 GGUF id, else None.""" + if is_unsloth_gemma4_edge_gguf(model_identifier): + return _GEMMA4_EDGE_TEMPLATE_FILE + if is_unsloth_gemma4_gguf(model_identifier): + return _GEMMA4_TEMPLATE_FILE + return None + + +@lru_cache(maxsize=8) +def load_bundled_chat_template(name: str) -> str: + """Read a bundled chat-template asset by filename (cached for the process).""" + return (_ASSETS_DIR / name).read_text(encoding="utf-8") + + +def resolve_effective_chat_template_override( + *, + model_identifier: Optional[str], + user_override: Optional[str], +) -> Optional[str]: + """Resolve which chat-template text to launch llama-server with. + + Precedence: + 1. An explicit, non-empty user override always wins (advanced users). + 2. For ``unsloth/gemma-4-*-GGUF``, return the bundled gemma-4 template + (adds ``preserve_thinking``, default off) so the embedded GGUF template + is overridden without re-downloading quants. E2B/E4B get the edge + variant; 12b/26B-A4B/31B get the standard one. + 3. Otherwise ``None`` -> llama-server renders the GGUF's embedded template. + + The result is fed to ``LlamaCppBackend.load_model(chat_template_override=...)`` + and must be computed before the route-level reload-dedup check so the live + backend state and the incoming request compare consistently. + """ + if user_override and user_override.strip(): + return user_override + template_file = _gemma4_template_file(model_identifier) + if template_file is not None: + return load_bundled_chat_template(template_file) + return None diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 55c2c551a3..2b9517692f 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -25,6 +25,7 @@ from utils.hardware import ( get_visible_gpu_count, ) from core.inference.audio_codecs import AudioCodecManager +from core.inference.runtime_context import runtime_context_length from io import StringIO import structlog from loggers import get_logger @@ -405,6 +406,10 @@ class InferenceBackend: # Reject CPU/disk offload for audio models too raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference") + self.models[model_name]["context_length"] = runtime_context_length( + self.models[model_name].get("model"), + max_seq_length, + ) self.active_model_name = model_name self.loading_models.discard(model_name) @@ -485,6 +490,10 @@ class InferenceBackend: self.models[model_name]["tokenizer"] = tokenizer raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference") + self.models[model_name]["context_length"] = runtime_context_length( + self.models[model_name].get("model"), + max_seq_length, + ) self._load_chat_template_info(model_name) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b83e4e6961..b5d5328fb0 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -28,10 +28,15 @@ from typing import Callable, Generator, Iterable, List, Optional import httpx from core.inference.llama_server_args import ( + extra_args_disable_mmproj, parse_cache_override, parse_ctx_override, + parse_split_mode_override, resolve_cache_type_kv, resolve_requested_ctx, + resolve_tensor_parallel, + strip_shadowing_flags, + strip_split_mode_only, ) from core.tool_healing import ( _TC_END_TAG_RE, @@ -59,10 +64,44 @@ from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, ) +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) logger = get_logger(__name__) +def _wsl_system_rocm_lib_dirs() -> "list[str]": + """System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL. + + The bundled bare-metal HIP can't drive WSL's /dev/dxg and segfaults on the + first GPU call; the system ROCm libs (libamdhip64 + librocdxg) can, while + the bundle still supplies libggml-hip / librocblas (gfx1151 kernels). + Mirrors install_llama_prebuilt._wsl_system_rocm_lib_dirs so a prebuilt that + passed install validation runs the same at serve time. No-op off a ROCDXG + WSL host (needs /dev/dxg, "microsoft" /proc/version, librocdxg in /opt/rocm). + """ + try: + if not os.path.exists("/dev/dxg"): + return [] + with open("/proc/version", encoding = "utf-8", errors = "replace") as fh: + if "microsoft" not in fh.read().lower(): + return [] + except OSError: + return [] + out: "list[str]" = [] + for d in ("/opt/rocm/lib", "/opt/rocm/lib64"): + if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists( + os.path.join(d, "librocdxg.so.1") + ): + out.append(d) + return out + + # ── Pre-compiled patterns for plan-without-action re-prompt ── # Forward-looking intent signals: the model is describing what it *will* # do rather than giving a final answer. @@ -683,6 +722,11 @@ class LlamaCppBackend: self._spec_fallback_reason: Optional[str] = None self._hf_variant: Optional[str] = None self._is_vision: bool = False + # Block-diffusion model (e.g. DiffusionGemma): served by the diffusion + # runner, not llama-server. Set from the GGUF architecture at load. + self._architecture: Optional[str] = None + self._is_diffusion: bool = False + self._diffusion_visual_bin: Optional[str] = None self._healthy = False # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None @@ -697,6 +741,8 @@ class LlamaCppBackend: self._supports_preserve_thinking: bool = False self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None + # Whether --split-mode tensor was applied on the active load. + self._tensor_parallel: bool = False self._reasoning_default: bool = True self._speculative_type: Optional[str] = None # Canonical UI-facing mode the user requested @@ -787,6 +833,11 @@ class LlamaCppBackend: def is_vision(self) -> bool: return self._is_vision + @property + def is_diffusion(self) -> bool: + """True when the loaded GGUF is a block-diffusion model (DiffusionGemma).""" + return self._is_diffusion + @property def hf_variant(self) -> Optional[str]: return self._hf_variant @@ -976,8 +1027,10 @@ class LlamaCppBackend: # enable_thinking / reasoning_effort -- skip. if self._supports_reasoning and not self._reasoning_always_on: if self._reasoning_style == "reasoning_effort": - if reasoning_effort in ("low", "medium", "high"): + if reasoning_effort in ("none", "low", "medium", "high"): kwargs["reasoning_effort"] = reasoning_effort + elif reasoning_effort == "minimal": + kwargs["reasoning_effort"] = "low" elif enable_thinking is not None: kwargs["reasoning_effort"] = "high" if enable_thinking else "low" else: @@ -995,6 +1048,11 @@ class LlamaCppBackend: def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv + @property + def tensor_parallel(self) -> bool: + """Whether --split-mode tensor is active on the loaded server.""" + return self._tensor_parallel + @property def speculative_type(self) -> Optional[str]: return self._speculative_type @@ -1321,6 +1379,105 @@ class LlamaCppBackend: return False return False + # Datacenter / professional NVIDIA parts that benefit from the llama.cpp + # FP32-accum / P2P tunings. Whole-word (\b) so short markers don't match + # workstation parts as substrings: "a100" must not fire on "RTX A1000". + _DATACENTER_GPU_RE = re.compile( + r"\b(?:a100|a30|h100|h200|h800|gh200|b200|b100|b300|gb200|gb300|" + r"l40s?|l4|rtx pro 6000|rtx 6000 ada)\b" + ) + + @staticmethod + def _is_datacenter_gpu(gpu_indices = None) -> bool: + """True iff every selected NVIDIA GPU is a datacenter/professional part. + NVIDIA-only, fails open to False (consumer GeForce, ROCm, CPU and errors + are left untouched); a mixed DC+consumer selection counts as non-DC. + + gpu_indices are PHYSICAL ids (see _get_gpu_free_memory), but + get_device_properties wants mask-relative ordinals, so we rebuild the + ordinal->physical map from CUDA_VISIBLE_DEVICES and key names by physical + id. Otherwise a masked host (CUDA_VISIBLE_DEVICES=4,5,6,7, selection [4,5]) + would drop the tuning or probe the wrong GPU.""" + try: + import torch + + if getattr(torch.version, "hip", None) is not None: + return False # ROCm reuses torch.cuda.*; not a CUDA part + if not (hasattr(torch, "cuda") and torch.cuda.is_available()): + return False + count = torch.cuda.device_count() + + # Mirror _get_gpu_free_memory: map visible ordinal -> physical id via + # CUDA_VISIBLE_DEVICES; unset/unparsable leaves physical id == ordinal. + physical_ids: Optional[list[int]] = None + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None: + try: + physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + physical_ids = None + + pattern = LlamaCppBackend._DATACENTER_GPU_RE + names_by_id: dict[int, str] = {} + for ordinal in range(count): + try: + name = (torch.cuda.get_device_properties(ordinal).name or "").lower() + except Exception: + continue + pid = ( + physical_ids[ordinal] + if physical_ids is not None and ordinal < len(physical_ids) + else ordinal + ) + names_by_id[pid] = name + + indices = list(gpu_indices) if gpu_indices else list(names_by_id) + saw = False + for _i in indices: + name = names_by_id.get(_i) + if name is None: + continue # not visible -> skip (fail conservative) + saw = True + if not pattern.search(name): + return False + return saw + except Exception: + return False + + @staticmethod + def _effective_gpu_count(gpu_indices = None) -> int: + """GPUs llama-server will use: len(selection), else the visible CUDA + device count (None = every visible GPU). 0 on error so multi-GPU tuning + stays off when the count is unknown.""" + if gpu_indices is not None: + return len(gpu_indices) + try: + import torch + if hasattr(torch, "cuda") and torch.cuda.is_available(): + return torch.cuda.device_count() + except Exception: + return 0 + return 0 + + @staticmethod + def _apply_datacenter_env(env: dict, gpu_indices = None) -> bool: + """Inject DC llama.cpp tuning into env in place via setdefault (user + values win); return whether the box qualified. Opt out with + UNSLOTH_DISABLE_DC_TUNING=1; only datacenter NVIDIA parts qualify + (consumer/ROCm/CPU/error are a no-op). Sets GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F + for any qualifying GPU (FP32 accum: ~0% cost on B200, real cost on GeForce), + plus GGML_CUDA_P2P + CUDA_SCALE_LAUNCH_QUEUES=4x for multi-GPU (+33-51% pp + tensor-split, +8-16% pipeline split on B200).""" + if os.environ.get("UNSLOTH_DISABLE_DC_TUNING") == "1": + return False + if not LlamaCppBackend._is_datacenter_gpu(gpu_indices): + return False + env.setdefault("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F", "1") + if LlamaCppBackend._effective_gpu_count(gpu_indices) > 1: + env.setdefault("GGML_CUDA_P2P", "1") + env.setdefault("CUDA_SCALE_LAUNCH_QUEUES", "4x") + return True + @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: """Query free memory per GPU. @@ -1499,6 +1656,23 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 + # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the + # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes + # via graph_reserve -- it is roughly EQUAL on every device (not proportional + # to the tensor split) and independent of context. Measured ~2.3 GB + # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a + # conservative headroom above that. It is (a) subtracted from each GPU's free + # VRAM before computing --tensor-split, so the roomier GPU absorbs more + # weight and the smallest GPU keeps room for KV, and (b) reserved per device + # when capping context. The auto-fallback to layer split covers any + # underestimate. NOTE: scales with the model's vocab / batch size; tune if a + # large-vocab model OOMs at load. + _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache + # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. + _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -1801,6 +1975,7 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + budget_frac: Optional[float] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. @@ -1834,8 +2009,11 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) + # MTP engaged: carve the drafter's reserve out of the fit budget. Callers + # can override outright (tensor-parallel mode passes a fatter margin), so + # only compute a default when none was supplied. + if budget_frac is None: + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) budget_bytes = available_mib * 1024 * 1024 * budget_frac model_footprint = model_size_bytes @@ -2050,11 +2228,16 @@ class LlamaCppBackend: self._ssm_state_size = None self._shared_kv_layers = None self._nextn_predict_layers = None + self._architecture = None + self._is_diffusion = False try: + canvas_seen = False WANTED = { "general.architecture", "tokenizer.chat_template", + # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. + "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. "general.source.huggingface.repository", "general.source.url", @@ -2109,6 +2292,7 @@ class LlamaCppBackend: general[key] = val_s if key == "general.architecture": arch = val_s + self._architecture = val_s arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", @@ -2137,6 +2321,8 @@ class LlamaCppBackend: if vtype == 4 else struct.unpack(" Optional[tuple[list, str, Optional[str]]]: + """Resolve how to launch the DiffusionGemma runner: (shim argv prefix, + visual-server binary, optional extra PYTHONPATH dir for the file override). + + Shim: UNSLOTH_DG_SHIM (a .py file) first, else the installed + unsloth_zoo.diffusion_studio.shim. Binary: DG_VISUAL_BIN first, else + alongside llama-server. Returns None if neither can be found. + """ + import importlib.util + import os + import sys + + # Visual-server binary: env override, else next to llama-server or in the + # install's build/bin (where the prebuilt/installer puts it). .exe on Windows. + visual_bin = os.environ.get("DG_VISUAL_BIN") + if not visual_bin: + name = "llama-diffusion-gemma-visual-server" + (".exe" if os.name == "nt" else "") + base = self._find_llama_server_binary() + if base: + base_dir = Path(base).parent + for cand in ( + base_dir / name, + base_dir / "build" / "bin" / name, + base_dir / "build" / "bin" / "Release" / name, + ): + if cand.is_file(): + visual_bin = str(cand) + break + if not (visual_bin and Path(visual_bin).is_file()): + return None + + # Shim: a file override (its dir goes on PYTHONPATH), else the zoo package via -m. + shim_file = os.environ.get("UNSLOTH_DG_SHIM") + if shim_file and Path(shim_file).is_file(): + return ([sys.executable, shim_file], visual_bin, str(Path(shim_file).parent)) + + # Find the installed shim without importing the heavy unsloth_zoo package + # (find_spec on the top-level package does not run its __init__). + try: + spec = importlib.util.find_spec("unsloth_zoo") + except Exception: + spec = None + if spec is not None and spec.submodule_search_locations: + pkg_dir = Path(list(spec.submodule_search_locations)[0]) + if (pkg_dir / "diffusion_studio" / "shim.py").is_file(): + return ( + [sys.executable, "-m", "unsloth_zoo.diffusion_studio.shim"], + visual_bin, + None, + ) + + return None + + def _start_diffusion_server( + self, + *, + model_path: str, + gguf_path: Optional[str], + hf_repo: Optional[str], + hf_variant: Optional[str], + model_identifier: str, + n_ctx: int, + extra_args: Optional[List[str]], + ) -> bool: + """Launch the OpenAI-compat diffusion shim (which drives the on-device + visual decoder) and wait for health. Presents the same /v1 + /health + interface as llama-server, so the rest of Studio is unchanged. + """ + import os + + assets = self._find_diffusion_assets() + if assets is None: + raise RuntimeError( + "DiffusionGemma runner not found. Install unsloth_zoo (which ships " + "unsloth_zoo.diffusion_studio.shim) or set UNSLOTH_DG_SHIM to a shim " + "file, and provide the visual-server binary via DG_VISUAL_BIN or next " + "to llama-server in the install tree." + ) + shim_cmd, visual_bin, extra_pythonpath = assets + self._diffusion_visual_bin = visual_bin + + self._kill_process() + self._port = self._find_free_port() + # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM + # (capped at the training context). An explicit in-range n_ctx overrides it. + maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 + gpu = os.environ.get("DG_GPU", "0") + + cmd = list(shim_cmd) + [ + "--gguf", + model_path, + "--host", + "127.0.0.1", + "--port", + str(self._port), + "--gpu", + gpu, + "--maxtok", + str(maxtok), + ] + + env = child_env_without_native_path_secret() + env["DG_VISUAL_BIN"] = visual_bin + env["DG_GPU"] = gpu + # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. + # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) + if extra_pythonpath: + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + (extra_pythonpath + os.pathsep + existing) if existing else extra_pythonpath + ) + + logger.info(f"Starting DiffusionGemma runner: {' '.join(cmd)}") + self._stdout_lines = [] + self._llama_log_fh = None + self._llama_log_path = None + try: + log_dir = _swa_cache_path().parent / "logs" / "diffusion-server" + log_dir.mkdir(parents = True, exist_ok = True) + self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log" + self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1) + logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}") + except OSError as e: + logger.debug(f"Could not open diffusion runner log file: {e}") + + # PR_SET_PDEATHSIG: the shim (and its visual server) die with this backend + # process, so a Studio crash/restart never orphans a GPU process. + popen_kwargs = dict(_windows_hidden_subprocess_kwargs()) + if sys.platform.startswith("linux"): # prctl/libc.so.6 are Linux-only + + def _pdeathsig(): + try: + import ctypes + import signal as _signal + ctypes.CDLL("libc.so.6", use_errno = True).prctl(1, _signal.SIGTERM) + except Exception: + pass + + popen_kwargs["preexec_fn"] = _pdeathsig + + self._process = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **popen_kwargs, + ) + self._stdout_thread = threading.Thread( + target = self._drain_stdout, daemon = True, name = "diffusion-stdout" + ) + self._stdout_thread.start() + + # Publish state before the health wait (mirrors the llama-server path). + self._gguf_path = model_path + self._hf_repo = hf_repo + self._is_vision = False + self._is_audio = False # clear any prior TTS/audio model's routing flag + self._model_identifier = model_identifier + self._cache_type_kv = None + self._gpu_offload_active = True + if hf_variant: + self._hf_variant = hf_variant + elif gguf_path: + try: + from utils.models.model_config import _extract_quant_label + self._hf_variant = _extract_quant_label(gguf_path) + except Exception: + self._hf_variant = None + else: + self._hf_variant = None + # Provisional until the server reports the budget it resolved (auto-size picks it from VRAM). + self._effective_context_length = maxtok or self._context_length + self._max_context_length = self._context_length or maxtok or None + + healthy = self._wait_for_health(timeout = 600.0) + if healthy: + self._healthy = True + self._gpu_offload_active = True + if extra_args is not None: + self._extra_args = list(extra_args) + self._extra_args_source = (model_identifier, hf_variant) + # The visual server logs "MAXTOK=" with the context budget it actually resolved + # (auto-sized to VRAM). Read it back so the UI context bar shows the real budget. + chosen = maxtok + try: + import re as _re + for _ln in reversed(self._stdout_lines): + _m = _re.search(r"MAXTOK=(\d+)", _ln) + if _m: + chosen = int(_m.group(1)) + break + except Exception: + pass + if chosen and chosen > 0: + self._effective_context_length = chosen + self._max_context_length = chosen + self._requested_n_ctx = int(n_ctx) + else: + self._healthy = False + logger.error("DiffusionGemma runner failed to become healthy") + return healthy + # ── HF download (no lock held) ─────────────────────────────── def _download_gguf( @@ -2643,6 +3046,16 @@ class LlamaCppBackend: """ lowered = (output or "").lower() + # Tensor parallelism (--split-mode tensor) is arch-gated in llama.cpp; + # unsupported architectures abort the load with this marker. Point the + # user at the toggle instead of a generic invalid-GGUF/OOM message. + if "split_mode_tensor not implemented" in lowered: + return ( + "Tensor parallelism is not supported for this model's " + "architecture. Turn off Tensor Parallelism in the model " + "settings and reload." + ) + # Detect Ollama source up front so the arch branch can keep the # Ollama hint instead of the generic "unsupported arch" message. gguf = gguf_path or "" @@ -2699,6 +3112,97 @@ class LlamaCppBackend: "Check that the GGUF file is valid and you have enough memory." ) + def _plan_tensor_parallel( + self, + gpus: list[tuple[int, int]], + model_size: int, + target_ctx: int, + cache_type_kv: Optional[str] = None, + n_parallel: int = 1, + mtp_engaged: bool = False, + max_target_ctx: Optional[int] = None, + ) -> tuple[int, int, list[int], Optional[list[int]]]: + """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. + + ``gpus`` is a list of ``(gpu_index, free_mib)``; ``model_size`` is the + weight size in bytes; ``target_ctx`` is the context to fit (the explicit + request, or the model's native length for auto). ``max_target_ctx`` is + the native/hardware ceiling used only for the UI bound (defaults to + ``target_ctx``). Returns + ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. + + Policy (assumes >= 2 GPUs; the caller drops the toggle below that): + - Cap context to the KV that fits the pooled VRAM after the weights and + one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only + cap, honored even for an explicit ``-c``. It is more accurate than the + 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. + - ``tensor_split`` is None (llama.cpp's even default, safe for every arch + incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even + share fits the smallest GPU; otherwise it is weighted by + ``(free - buffer)`` so the roomier GPU absorbs more weight and the + smallest GPU keeps room for KV. + """ + # Drop GPUs that can't hold the per-device compute-graph buffer; they'd + # OOM in tensor mode. load_model already filters before calling, so this + # is defense-in-depth that also keeps the pure function self-contained + # (and unit-testable without a GPU). + reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + gpu_indices = sorted(idx for idx, _ in usable_gpus) + if len(gpu_indices) < 2: + # Tensor parallelism is meaningless on <2 GPUs (the caller drops the + # toggle before this); be defensive and never emit a split here. + return ( + target_ctx if target_ctx > 0 else 4096, + target_ctx if target_ctx > 0 else 4096, + gpu_indices, + None, + ) + free_by_idx = {idx: free for idx, free in usable_gpus} + pool_mib = sum(free_by_idx.values()) + kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size + if mtp_engaged: + # MTP keeps a draft model + its own KV cache on GPU. + kv_budget_b -= 2 * 1024**3 + + def _fit_ctx(ctx: int) -> int: + # Largest context whose KV fits the pooled budget. Floors small, but + # never raises an explicit ctx above what was asked. + if self._can_estimate_kv() and ctx > 0: + ctx_floor = min(2048, ctx) + if kv_budget_b <= 0: + # Weights + buffers exceed the pool -> floor; the load then + # falls back to layer split. + return ctx_floor + kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) + if kv_at <= kv_budget_b: + return ctx + return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + # KV size unknown -> can't prove a safe cap; floor. + return min(4096, ctx) if ctx > 0 else 4096 + + # max_available_ctx is the hardware ceiling for the UI bound, sized from + # the native context independent of an explicit small -c (which only + # caps effective_ctx). + max_ctx_target = max_target_ctx if (max_target_ctx and max_target_ctx > 0) else target_ctx + max_available_ctx = _fit_ctx(max_ctx_target) + effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) + + min_free_mib = min(free_by_idx.values()) + kv_bytes = ( + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) + if (self._can_estimate_kv() and effective_ctx > 0) + else 0 + ) + even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + tensor_split: Optional[list[int]] = None + if even_share_mib > (min_free_mib - reserve_mib): + adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if sum(adj) > 0: + tensor_split = adj + return effective_ctx, max_available_ctx, gpu_indices, tensor_split + @staticmethod def _is_projector_incompatibility(output: str) -> bool: """True when llama-server aborted because it cannot load the model's @@ -2821,6 +3325,7 @@ class LlamaCppBackend: cache_type_kv: Optional[str] = None, speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, + tensor_parallel: bool = False, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -2852,6 +3357,7 @@ class LlamaCppBackend: cache_type_kv = cache_type_kv, speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, + tensor_parallel = tensor_parallel, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -2908,13 +3414,9 @@ class LlamaCppBackend: with self._lock: self._kill_process() + # Resolve llama-server now but defer a not-found error: a block-diffusion + # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() - if not binary: - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -2929,8 +3431,8 @@ class LlamaCppBackend: hf_variant = hf_variant, hf_token = hf_token, ) - # Auto-download mmproj for vision models - if is_vision and not mmproj_path: + # Auto-download mmproj for vision models unless opted out. + if is_vision and not mmproj_path and not extra_args_disable_mmproj(extra_args): mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, @@ -2969,6 +3471,30 @@ class LlamaCppBackend: logger.info("Load cancelled after download phase") return False + # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; + # serve them with the diffusion runner (same OpenAI-compat interface). + if self._is_diffusion: + with self._lock: + if self._cancel_event.is_set(): + logger.info("Load cancelled before diffusion server start") + return False + return self._start_diffusion_server( + model_path = model_path, + gguf_path = gguf_path, + hf_repo = hf_repo, + hf_variant = hf_variant, + model_identifier = model_identifier, + n_ctx = n_ctx, + extra_args = extra_args, + ) + + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the # frontend /unload+/load Apply path engages the wait here even @@ -2991,10 +3517,43 @@ class LlamaCppBackend: requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) + # A user --split-mode in extras last-wins-overrides the + # toggle, so reconcile it back into tensor_parallel state. + split_mode_override = parse_split_mode_override(extra_args) + tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + # Tensor mode aborts on a quantized KV cache, so drop it for the + # tensor attempt (and strip any inherited/explicit --cache-type + # that would re-impose it when appended last). The layer-split + # fallback re-runs with tensor_parallel False and keeps the type. + if ( + tensor_parallel + and cache_type_kv + and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + ): + logger.info( + "Tensor parallelism requires a non-quantized KV cache; " + "ignoring cache type %s for the tensor attempt.", + cache_type_kv, + ) + cache_type_kv = None + if extra_args: + extra_args = strip_shadowing_flags( + extra_args, + strip_context = False, + strip_cache = True, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + if split_mode_override is not None: + logger.info( + f"User --split-mode {split_mode_override} honored; " + "reconciled into tensor_parallel state" + ) effective_ctx = requested_ctx if requested_ctx > 0 else (self._context_length or 0) max_available_ctx = self._context_length or effective_ctx gpus: list[tuple[int, int]] = [] @@ -3056,6 +3615,8 @@ class LlamaCppBackend: # Auto n_ctx=0 (native): prefer fewer GPUs with reduced # context, since multi-GPU is slower. gpu_indices, use_fit = None, True + # Per-GPU weight proportions for tensor mode (None = even). + tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # MTP draft model lives outside the main estimates; carve # its reserve out of every fit budget and pin threshold so @@ -3063,10 +3624,67 @@ class LlamaCppBackend: _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve - if gpus and self._can_estimate_kv() and effective_ctx > 0: - # Largest hardware-aware cap from the native context - # across all usable GPU subsets (for UI bounds), - # independent of the requested context. + # Tensor mode allocates a compute-graph buffer on every + # participating GPU, so a GPU with less free VRAM than that + # reserve can't host it and would OOM at load. Drop those + # from the tensor-parallel set up front (gpu_indices below + # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded + # from llama-server entirely, not just given zero weight). + tp_gpus = gpus + if tensor_parallel: + reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + + if tensor_parallel and len(tp_gpus) < 2: + # Tensor parallelism needs >= 2 usable GPUs. On a single + # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only + # or probe failed) it must not reach llama-server; and a + # GPU below the buffer reserve can't participate. Drop the + # flag and fall through to normal layer/CPU allocation. + logger.info( + "Tensor parallelism requested but only %d of %d GPU(s) " + "have enough free VRAM for the compute buffer; " + "ignoring (needs >= 2).", + len(tp_gpus), + len(gpus), + ) + tensor_parallel = False + # A user --split-mode tensor in extras is appended after + # Studio's flags, so it would still reach llama-server and + # fail here; strip it so the downgrade actually applies. + extra_args = strip_split_mode_only(extra_args) + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation: use all usable GPUs, weight + # the split by (free - buffer), and cap context to the + # pooled VRAM after weights + per-device compute-graph + # buffers. See _plan_tensor_parallel for the policy. + target_ctx = ( + effective_ctx + if explicit_ctx + else (self._context_length or effective_ctx) + ) + ( + effective_ctx, + max_available_ctx, + gpu_indices, + tp_tensor_split, + ) = self._plan_tensor_parallel( + tp_gpus, + model_size, + target_ctx, + cache_type_kv = cache_type_kv, + n_parallel = n_parallel, + mtp_engaged = _mtp_will_engage, + # Report the UI ceiling from native ctx, not the + # explicit small request. + max_target_ctx = self._context_length or target_ctx, + ) + use_fit = False + elif gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) @@ -3189,12 +3807,15 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"GPU selection failed ({e}), using --fit on") gpu_indices, use_fit = None, True + tp_tensor_split = None effective_ctx = requested_ctx # fall back to original - launch_mmproj_path = self._resolve_launch_mmproj_path( - model_path = model_path, - mmproj_path = mmproj_path, - ) + launch_mmproj_path = None + if not extra_args_disable_mmproj(extra_args): + launch_mmproj_path = self._resolve_launch_mmproj_path( + model_path = model_path, + mmproj_path = mmproj_path, + ) # Need both a resolved mmproj AND the config vision flag; a stray # mmproj passing the family-name heuristic must not flip a non-VLM # GGUF into vision mode. @@ -3286,6 +3907,27 @@ class LlamaCppBackend: else: self._cache_type_kv = None + # Tensor parallelism: split the model across GPUs by tensor + # rather than by layer. Multi-GPU only -- a no-op on a single + # GPU. Default (layer split) is left implicit by omitting the + # flag. See llama.cpp --split-mode. + if tensor_parallel: + cmd.extend(["--split-mode", "tensor"]) + if tp_tensor_split and len(tp_tensor_split) > 1: + cmd.extend( + [ + "--tensor-split", + ",".join(str(int(x)) for x in tp_tensor_split), + ] + ) + self._tensor_parallel = True + logger.info( + "Tensor parallelism: --split-mode tensor, --tensor-split %s", + tp_tensor_split, + ) + else: + self._tensor_parallel = False + # Speculative decoding. See _build_speculative_flags for the # mode resolution, benchmarks, and llama.cpp references. launch_mtp_draft_path = self._resolve_launch_mtp_path( @@ -3324,6 +3966,7 @@ class LlamaCppBackend: self._chat_template_file = tempfile.NamedTemporaryFile( mode = "w", + encoding = "utf-8", suffix = ".jinja", delete = False, prefix = "unsloth_chat_template_", @@ -3345,6 +3988,13 @@ class LlamaCppBackend: thinking_default = False self._reasoning_default = thinking_default reasoning_kw = self._reasoning_kwargs(thinking_default) + # preserve_thinking is an independent kwarg. Default it OFF + # at launch so direct OpenAI-compatible callers that omit the + # field match the UI's default-off behavior (the bundled + # gemma-4 template also defaults it false; the frontend sends + # preserve_thinking per request once toggled on). + if self._supports_preserve_thinking: + reasoning_kw["preserve_thinking"] = False cmd.extend( [ "--chat-template-kwargs", @@ -3395,6 +4045,14 @@ class LlamaCppBackend: env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") + # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). + # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. + if self._apply_datacenter_env(env, gpu_indices): + multi_gpu = self._effective_gpu_count(gpu_indices) > 1 + logger.info( + f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" + ) + if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. path_dirs = self._build_windows_path_dirs( @@ -3420,7 +4078,15 @@ class LlamaCppBackend: # plus CUDA runtime libs (libcudart, libcublas, etc.) import platform - lib_dirs = [binary_dir] + lib_dirs = [] + # WSL: system HIP before the bundle's (which segfaults on + # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which + # validates the prebuilt with this same ordering. + for _wsl_rocm in _wsl_system_rocm_lib_dirs(): + lib_dirs.append(_wsl_rocm) + if lib_dirs: + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + lib_dirs.append(binary_dir) _arch = platform.machine() # x86_64, aarch64, etc. # Pip-installed nvidia CUDA runtime libs. The prebuilt @@ -3745,7 +4411,7 @@ class LlamaCppBackend: ) logger.info( - f"llama-server ready on port {self._port} " f"for model '{model_identifier}'" + f"llama-server ready on port {self._port} for model '{model_identifier}'" ) # Probe outside _lock (interruptible by /unload); init inside. @@ -4054,6 +4720,7 @@ class LlamaCppBackend: is_vision: bool, gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, + tensor_parallel: bool = False, mtp_draft_path: Optional[str] = None, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -4090,6 +4757,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False + # Reconcile a user --split-mode in extras (load_model does the same), so + # an extras-driven tensor load isn't seen as a mismatch that needlessly + # kills/reloads a healthy server. + if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -4161,6 +4834,12 @@ class LlamaCppBackend: return None return saw_gpu_buffer + def load_cancelled(self) -> bool: + """True if a load was cancelled (e.g. via unload/_cancel_event) and not + yet consumed by the next load_model. Lets the tensor->layer fallback + avoid restarting a load the user just cancelled.""" + return self._cancel_event.is_set() + def unload_model(self) -> bool: """Terminate the subprocess and cancel any in-flight download.""" self._cancel_event.set() @@ -4193,6 +4872,7 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + self._tensor_parallel = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None @@ -4370,7 +5050,7 @@ class LlamaCppBackend: proc.kill() logger.info( - f"Killed orphaned llama-server process " f"(pid={proc.info['pid']})" + f"Killed orphaned llama-server process (pid={proc.info['pid']})" ) except ( psutil.NoSuchProcess, @@ -4823,6 +5503,12 @@ class LlamaCppBackend: try: data = json.loads(line[6:]) + # Diffusion frame (per-step canvas) from the shim: forward untouched so + # the frontend renders it in place. No assistant text, so it never enters + # the cumulative content. + if data.get("type") == "diffusion_frame": + yield data + continue # Capture server timings/usage from final chunks. _chunk_timings = data.get("timings") if _chunk_timings: @@ -4907,6 +5593,7 @@ class LlamaCppBackend: rag_scope: Optional[dict] = None, seed: Optional[int] = None, disable_parallel_tool_use: bool = False, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -4925,7 +5612,7 @@ class LlamaCppBackend: # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits # the same tool card + citations a real call would. - _auto = build_rag_autoinject(conversation, rag_scope) + _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -5074,7 +5761,7 @@ class LlamaCppBackend: if response.status_code != 200: error_body = response.read().decode() raise RuntimeError( - f"llama-server returned {response.status_code}: " f"{error_body}" + f"llama-server returned {response.status_code}: {error_body}" ) raw_buf = "" @@ -5485,8 +6172,7 @@ class LlamaCppBackend: force = True, ) logger.info( - f"Safety net: parsed {len(tool_calls)} tool call(s) " - f"from streamed content" + f"Safety net: parsed {len(tool_calls)} tool call(s) from streamed content" ) else: # ── DRAINING path: assemble tool_calls ── @@ -5606,8 +6292,51 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - yield {"type": "status", "text": decision.status_text} - yield decision.tool_start_event() + needs_confirm = bool(confirm_tool_calls) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = ( + begin_tool_decision(session_id, approval_id) if needs_confirm else None + ) + start_event = decision.tool_start_event() + start_event["approval_id"] = approval_id + start_event["awaiting_confirmation"] = needs_confirm + + try: + yield {"type": "status", "text": decision.status_text} + yield start_event + + if ( + decision_slot is not None + and wait_tool_decision( + decision_slot, + approval_id, + cancel_event = cancel_event, + ) + == "deny" + ): + decision_slot = None + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": TOOL_REJECTED_MESSAGE, + "provenance": decision.provenance, + } + denied_message = { + "role": "tool", + "name": decision.tool_name, + "content": TOOL_REJECTED_MESSAGE, + } + if decision.tool_call_id: + denied_message["tool_call_id"] = decision.tool_call_id + conversation.append(denied_message) + if _forced_tool_call_pending: + _forced_tool_call_pending = False + continue + decision_slot = None + finally: + if decision_slot is not None: + abort_tool_decision(decision_slot, approval_id) _effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout # RAG: cap paraphrased KB re-searches that slip past the dup guard. diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 1862c9c5de..69a86fa3ba 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -109,6 +109,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: out.append(token) parse_ctx_override(out) parse_cache_override(out) + parse_split_mode_override(out) return out @@ -157,8 +158,20 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset( "--no-jinja", } ) +# Multi-GPU split mode shadows the Tensor Parallelism toggle +# (--split-mode tensor). Pass-through stays allowed so users keep the +# row/none/layer modes the toggle doesn't expose, but it's stripped on +# inherit and reconciled into the round-tripped tensor_parallel state. +# --tensor-split is coupled to the split mode and is stripped with it: Studio +# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must +# not last-wins-override Studio's computed asymmetric split. +_SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"}) +_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) +_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS -_SHADOWING_FLAGS: frozenset[str] = _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS +_SHADOWING_FLAGS: frozenset[str] = ( + _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS +) # Shadowing flags that take no value -- strip the flag only, not the next token. _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) @@ -213,11 +226,12 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> return override if override is not None else fallback_n_ctx -def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: - """Return the last-wins cache type if extras pass cache flags. +def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]: + """Return the last-wins string value among ``flags`` in extras, or None. - Recognises -ctk (key) and -ctv (value); treats both as one setting, - since Studio's KV estimate has a single cache_type_kv knob. + Handles both ``--flag=value`` and ``--flag value`` forms and raises if a + matched flag has no (or an empty) value. Shared by the single-knob + last-wins parsers (cache type, split mode). """ if not args: return None @@ -228,7 +242,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: while i < n: tok = tokens[i] flag = _flag_name(tok) - if flag is None or flag not in _CACHE_FLAGS: + if flag is None or flag not in flags: i += 1 continue @@ -249,6 +263,17 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return override +def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins cache type if extras pass cache flags. + + Mirrors parse_ctx_override but for cache type. Recognises both -ctk + (key) and -ctv (value). When both flags appear, returns the last-wins + value, treating key and value cache flags as the same setting because + Studio's KV estimate has a single cache_type_kv knob. + """ + return _last_flag_value(args, _CACHE_FLAGS) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -260,6 +285,52 @@ def resolve_cache_type_kv( return override if override is not None else fallback_cache_type_kv +def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins ``--split-mode`` / ``-sm`` value from extras. + + Mirrors parse_cache_override for the multi-GPU split mode. Returns the + raw mode string (e.g. ``tensor`` / ``row`` / ``none`` / ``layer``), or + None when extras don't set it. + """ + return _last_flag_value(args, _SPLIT_MODE_FLAGS) + + +def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool: + """Return the tensor-parallel state load_model should treat as requested. + + A user-supplied ``--split-mode`` in extras last-wins-overrides the + toggle, so reconcile it back into the boolean: any explicit split mode + means tensor-parallel is on iff that mode is ``tensor``. Falls back to + the toggle value when extras don't set it. + """ + override = parse_split_mode_override(args) + if override is None: + return fallback_tensor_parallel + return override.strip().lower() == "tensor" + + +_MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) +_MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) + + +def extra_args_disable_mmproj(args: Optional[Iterable[str]]) -> bool: + """True when pass-through args opt out of vision mmproj loading. + + llama-server parses --mmproj-auto / --no-mmproj / --no-mmproj-auto as one + boolean with last-wins semantics; mirror that here. + """ + if not args: + return False + disabled = False + for raw in args: + flag = _flag_name(str(raw)) + if flag in _MMPROJ_DISABLE_FLAGS: + disabled = True + elif flag in _MMPROJ_ENABLE_FLAGS: + disabled = False + return disabled + + def strip_shadowing_flags( args: Iterable[str], *, @@ -267,12 +338,15 @@ def strip_shadowing_flags( strip_cache: bool = True, strip_spec: bool = True, strip_template: bool = True, + strip_split_mode: bool = True, ) -> list[str]: """Strip flags that shadow first-class Studio settings. Used when inheriting a previous load's ``llama_extra_args`` so an - inherited `-c 4096` can't override the current `max_seq_length` (same for - cache / spec / template). Each ``strip_*`` toggle controls one group. + inherited `-c 4096` can't override the current `max_seq_length` + (same for cache / spec / template / split-mode). Each ``strip_*`` + toggle controls one group; the route only strips groups whose + first-class field the caller actually supplied. """ shadowing: set[str] = set() if strip_context: @@ -283,6 +357,8 @@ def strip_shadowing_flags( shadowing |= _SPEC_FLAGS if strip_template: shadowing |= _TEMPLATE_FLAGS + if strip_split_mode: + shadowing |= _SPLIT_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] @@ -303,3 +379,20 @@ def strip_shadowing_flags( else: i += 1 return out + + +def strip_split_mode_only(args: Optional[Iterable[str]]) -> Optional[list[str]]: + """Remove the split-mode group (``--split-mode`` / ``-sm`` and the coupled + ``--tensor-split`` / ``-ts``) from ``args``, keeping every other shadow flag. + Preserves a None/empty input so the inherit-vs-explicit-empty distinction + survives. Used where tensor mode is being forced off (downgrade / fallback).""" + if not args: + return args + return strip_shadowing_flags( + args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index a0d79bbdf4..5a36d90c5d 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -8,6 +8,7 @@ import json import os import shlex import sys +import time from typing import Any, Optional from loggers import get_logger @@ -16,6 +17,14 @@ logger = get_logger(__name__) MCP_TOOL_PREFIX = "mcp__" +# A failed probe isn't cached (a recovered server must come back), but it's +# recorded so a down server isn't re-probed -- and the chat send re-hung for +# the full timeout -- on every message. Cool off for this long after a failure; +# much longer for OAuth, whose probe can hang up to _OAUTH_PROBE_TIMEOUT, +# so that hang doesn't recur every minute. +FAILED_PROBE_COOLOFF_SECONDS = 60.0 +OAUTH_FAILED_PROBE_COOLOFF_SECONDS = 300.0 + _oauth_token_store = None @@ -227,6 +236,53 @@ async def list_tools_async( return await asyncio.wait_for(_fetch(), timeout = timeout) +# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools() +# probes a server only on a cache miss, keeping MCP discovery off the chat +# send's critical path -- tool schemas are stable within a session. The +# /refresh route warms it; a URL/header/OAuth change or a delete evicts it. +# Successful probes are cached indefinitely. +_tool_cache: dict[str, list[dict]] = {} + +# server_id -> monotonic time before which a failed server must not be +# re-probed (see record_probe_failure). Cleared on a successful probe or +# eviction. +_probe_cooloff_until: dict[str, float] = {} + +# MCP server fields whose change invalidates a server's discovered tools: the +# endpoint/auth used to probe it (url, headers, oauth) or whether it's used at +# all (is_enabled). A rename does not. The update route's eviction and +# get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift. +TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"}) + + +def get_cached_tools(server_id: str) -> Optional[list[dict]]: + return _tool_cache.get(server_id) + + +def cache_tools(server_id: str, tools: list[dict]) -> None: + _tool_cache[server_id] = tools + _probe_cooloff_until.pop(server_id, None) + + +def record_probe_failure(server_id: str, use_oauth: bool = False) -> None: + cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS + _probe_cooloff_until[server_id] = time.monotonic() + cooloff + + +def in_failure_cooloff(server_id: str) -> bool: + return _probe_cooloff_until.get(server_id, 0.0) > time.monotonic() + + +def invalidate_tool_cache(server_id: Optional[str] = None) -> None: + """Evict one server's cached tools, or every entry when server_id is None.""" + if server_id is None: + _tool_cache.clear() + _probe_cooloff_until.clear() + else: + _tool_cache.pop(server_id, None) + _probe_cooloff_until.pop(server_id, None) + + def _flatten_result(result: Any) -> str: parts = [] for block in getattr(result, "content", None) or []: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index fbbdb9a773..5c7799152f 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -7,6 +7,7 @@ instead of torch/transformers for model loading and generation. import threading from typing import Optional, Generator +from core.inference.runtime_context import runtime_context_length from loggers import get_logger logger = get_logger(__name__) @@ -175,6 +176,7 @@ class MLXInferenceBackend: "is_audio": False, "audio_type": None, "has_audio_input": False, + "context_length": runtime_context_length(self._model, max_seq_length), } # Capture chat_template_info so the worker IPC reply ships it back and # the route layer classifies capabilities like the other paths. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 32bb25c976..4ccac2912e 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -727,6 +727,7 @@ class InferenceOrchestrator: "is_audio": model_info.get("is_audio", False), "audio_type": model_info.get("audio_type"), "has_audio_input": model_info.get("has_audio_input", False), + "context_length": model_info.get("context_length"), } # Mirror chat_template_info so routes can classify caps # without re-entering the subprocess. @@ -860,6 +861,7 @@ class InferenceOrchestrator: tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + confirm_tool_calls: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, **_unused, @@ -921,6 +923,7 @@ class InferenceOrchestrator: tool_call_timeout = tool_call_timeout, session_id = session_id, rag_scope = rag_scope, + confirm_tool_calls = confirm_tool_calls, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index f7a2da2833..5b72373c03 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -249,6 +249,23 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { # Surfaced via the frontend's CUSTOM_PROVIDER_PRESETS, not the dropdown. "hidden": True, }, + "custom": { + "display_name": "Custom", + # User-supplied via provider_base_url. + "base_url": "", + "default_models": [], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": ( + "User-supplied OpenAI-compatible server. Routed to " + "/v1/chat/completions; /models is optional." + ), + # Surfaced by the frontend's generic Custom option, not the dropdown. + "hidden": True, + }, "ollama": { "display_name": "Ollama", "base_url": "http://localhost:11434/v1", diff --git a/studio/backend/core/inference/runtime_context.py b/studio/backend/core/inference/runtime_context.py new file mode 100644 index 0000000000..8b881628be --- /dev/null +++ b/studio/backend/core/inference/runtime_context.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Runtime context length helpers shared by inference backends.""" + +from __future__ import annotations + +from typing import Any, Optional + + +def runtime_context_length(model: Any, fallback: Optional[int] = None) -> Optional[int]: + """Return the effective context length Unsloth attached to a loaded model.""" + for value in (getattr(model, "max_seq_length", None), fallback): + if isinstance(value, bool): + continue + try: + value_int = int(value) + except (TypeError, ValueError): + continue + if value_int > 0: + return value_int + return None diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 7942edb6d7..3b6a393f3d 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -35,6 +35,13 @@ from core.inference.tool_loop_controller import ( status_for_tool, tool_event_provenance, ) +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) logger = get_logger(__name__) @@ -146,6 +153,7 @@ def run_safetensors_tool_loop( tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -174,7 +182,7 @@ def run_safetensors_tool_loop( # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search. from core.inference.tools import build_rag_autoinject - _auto = build_rag_autoinject(conversation, rag_scope) + _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -509,8 +517,47 @@ def run_safetensors_tool_loop( else: assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) - yield {"type": "status", "text": decision.status_text} - yield decision.tool_start_event() + needs_confirm = bool(confirm_tool_calls) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None + start_event = decision.tool_start_event() + start_event["approval_id"] = approval_id + start_event["awaiting_confirmation"] = needs_confirm + + try: + yield {"type": "status", "text": decision.status_text} + yield start_event + + if ( + decision_slot is not None + and wait_tool_decision( + decision_slot, + approval_id, + cancel_event = cancel_event, + ) + == "deny" + ): + decision_slot = None + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": TOOL_REJECTED_MESSAGE, + "provenance": decision.provenance, + } + denied_message = { + "role": "tool", + "name": decision.tool_name, + "content": TOOL_REJECTED_MESSAGE, + } + if decision.tool_call_id: + denied_message["tool_call_id"] = decision.tool_call_id + conversation.append(denied_message) + continue + decision_slot = None + finally: + if decision_slot is not None: + abort_tool_decision(decision_slot, approval_id) eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout # RAG: cap paraphrased KB re-searches that slip past the dup guard. diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py new file mode 100644 index 0000000000..73687165b8 --- /dev/null +++ b/studio/backend/core/inference/tensor_fallback.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tensor-parallel -> layer-split auto-fallback for GGUF loads. + +Kept in its own module (no FastAPI / httpx deps) so the orchestration can be +unit-tested with a fake loader, without a GPU or a running llama-server. +""" + +from __future__ import annotations + +import logging +from typing import Awaitable, Callable, Optional + +from core.inference.llama_server_args import ( + resolve_tensor_parallel, + strip_split_mode_only, +) + +logger = logging.getLogger(__name__) + + +async def load_with_tensor_fallback( + attempt_load: Callable[[bool, Optional[list[str]]], Awaitable[bool]], + *, + requested_tensor: bool, + extra_args: Optional[list[str]], + label: str = "", + cancelled: Optional[Callable[[], bool]] = None, +) -> bool: + """Run a GGUF load with the tensor-parallel -> layer-split auto-fallback. + + ``attempt_load(tensor_parallel, extra_args)`` performs one load and returns + True on success; it *raises* on a hard crash (llama-server aborts on some + archs / older builds), which is treated the same as a False return. + + Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` + in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether + tensor mode is actually engaged, and it strips ``--split-mode`` from the + extras so the layer retry can't relaunch the same failing tensor load. A + non-tensor load keeps its original contract and propagates exceptions. + + ``cancelled()`` distinguishes a real tensor-start failure from a user + cancellation: ``attempt_load`` also returns False when the load was + cancelled, so without this the helper would restart a load the user just + cancelled. + """ + tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + try: + success = await attempt_load(requested_tensor, extra_args) + except Exception as exc: + if not tensor_requested: + raise + logger.warning("Tensor-parallel load raised for '%s': %s", label, exc) + success = False + + if success or not tensor_requested: + return success + + # The first attempt returned False because the user cancelled, not because + # tensor mode is unsupported -- do not relaunch the cancelled load. + if cancelled is not None and cancelled(): + return success + + logger.warning( + "Tensor-parallel load failed for '%s'; retrying with layer split " + "(this model may not support tensor parallelism)", + label, + ) + return await attempt_load(False, strip_split_mode_only(extra_args)) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index b27fa6ff73..8c45706395 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -24,11 +24,16 @@ import urllib.request from core.inference.mcp_client import ( MCP_TOOL_PREFIX, + TOOL_CACHE_INVALIDATING_FIELDS, + cache_tools, call_tool_sync, + get_cached_tools, + in_failure_cooloff, is_stdio, list_tools_async, parse_server_headers, probe_timeout, + record_probe_failure, stdio_mcp_enabled, ) from storage import mcp_servers_db @@ -634,28 +639,56 @@ async def get_enabled_mcp_tools() -> list[dict]: if not servers: return [] - results = await asyncio.gather( - *( - list_tools_async( - url = s["url"], - headers = parse_server_headers(s), - timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), - use_oauth = bool(s.get("use_oauth")), - ) - for s in servers - ), - return_exceptions = True, - ) + # Skip servers still in their post-failure cool-off, otherwise a down + # server gets re-probed -- and blocks the send for the full timeout -- on + # every message. + uncached = [ + s for s in servers if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"]) + ] + if uncached: + results = await asyncio.gather( + *( + list_tools_async( + url = s["url"], + headers = parse_server_headers(s), + timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), + use_oauth = bool(s.get("use_oauth")), + ) + for s in uncached + ), + return_exceptions = True, + ) + # An edit/delete can land while we await a probe (up to 305 s for + # OAuth); its cache eviction is a no-op against an entry we haven't + # written yet. Re-read and drop a result whose server changed or + # was removed mid-probe, else a stale tool list caches indefinitely. + current = {s["id"]: s for s in mcp_servers_db.list_servers()} + for server, payload in zip(uncached, results): + # Guard the failure branch too: a stale failure must not park a + # cool-off on the fresh config, or the server the user just fixed + # is skipped for the whole window. + fresh = current.get(server["id"]) + if fresh is None or any( + fresh.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + continue + if isinstance(payload, BaseException): + logger.warning( + "MCP server '%s' (%s) discovery failed: %s", + server.get("display_name") or server["id"], + server.get("url"), + payload, + ) + # Failures aren't cached, but record one so a down server + # isn't re-probed every send during the cool-off. + record_probe_failure(server["id"], bool(fresh.get("use_oauth"))) + continue + cache_tools(server["id"], payload) specs: list[dict] = [] - for server, payload in zip(servers, results): - if isinstance(payload, BaseException): - logger.warning( - "MCP server '%s' (%s) discovery failed: %s", - server.get("display_name") or server["id"], - server.get("url"), - payload, - ) + for server in servers: + payload = get_cached_tools(server["id"]) + if payload is None: continue specs.extend(_mcp_specs_for_server(server, payload)) return specs diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 4f5b985234..cc79654087 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -315,6 +315,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: "audio_type": getattr(mc, "audio_type", None), "has_audio_input": getattr(mc, "has_audio_input", False), } + try: + _bm = getattr(backend, "models", {}) or {} + _entry = ( + _bm.get(mc.identifier) + or _bm.get(getattr(backend, "active_model_name", None)) + or {} + ) + _context_length = _entry.get("context_length") + if _context_length is not None: + model_info["context_length"] = int(_context_length) + except Exception as _ctx_exc: + logger.warning("context_length forward failed: %s", _ctx_exc) # Forward chat_template_info so the parent can classify capabilities. try: _bm = getattr(backend, "models", {}) or {} @@ -881,6 +893,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf name: { "is_vision": info.get("is_vision", False), "is_lora": info.get("is_lora", False), + "context_length": info.get("context_length"), } for name, info in backend.models.items() }, diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index 165c1c2cf1..2a4a198610 100644 --- a/studio/backend/core/training/resume.py +++ b/studio/backend/core/training/resume.py @@ -3,6 +3,7 @@ """Helpers for validating resumable training outputs.""" +import json from pathlib import Path from typing import Optional @@ -56,9 +57,29 @@ def normalize_resume_output_dir(path_value: str) -> str: return str(path) +def _run_config(run: dict) -> dict: + raw_config = run.get("config_json") + if isinstance(raw_config, dict): + return raw_config + if not isinstance(raw_config, str) or not raw_config.strip(): + return {} + try: + parsed = json.loads(raw_config) + except (json.JSONDecodeError, TypeError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _uses_s3_dataset(run: dict) -> bool: + config = _run_config(run) + return config.get("dataset_source") == "s3" or "s3_dataset" in config + + def can_resume_run(run: dict) -> bool: if run.get("resumed_later"): return False + if _uses_s3_dataset(run): + return False final_step = run.get("final_step") total_steps = run.get("total_steps") diff --git a/studio/backend/core/training/s3_dataset.py b/studio/backend/core/training/s3_dataset.py new file mode 100644 index 0000000000..3d05d19c75 --- /dev/null +++ b/studio/backend/core/training/s3_dataset.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +S3 dataset loader. + +Downloads dataset files (parquet / json / jsonl / csv) from an AWS S3 bucket +to a local temp directory so the existing local-file dataset path can consume +them. boto3 is an optional dependency and is imported lazily — callers should +gate on :func:`boto3_available` before invoking the loader. + +The S3 config dict mirrors ``models.training.S3Config.model_dump()`` (snake_case +keys): bucket, region, prefix, access_key_id, secret_access_key, use_iam_role. +Credentials are read once to build the client and never logged or persisted. +""" + +from __future__ import annotations + +import logging +import os +import shutil +import tempfile +from importlib.util import find_spec +from typing import Callable, Optional + +logger = logging.getLogger(__name__) + +# Extensions the local-file loader (UnslothTrainer._loader_for_files) understands. +SUPPORTED_EXTENSIONS = (".parquet", ".json", ".jsonl", ".csv") +_JSON_EXTENSIONS = (".json", ".jsonl") +_IGNORED_METADATA_FILENAMES = { + "dataset_info.json", + "metadata.json", + "schema.json", + "state.json", +} + + +class S3DownloadCancelled(RuntimeError): + """Raised when the caller cancels an S3 dataset download.""" + + +class S3DatasetDownload: + def __init__( + self, + files: list[str], + temp_dir: Optional[str] = None, + ): + self.files = files + self.temp_dir = temp_dir + + def cleanup(self) -> None: + if not self.temp_dir: + return + shutil.rmtree(self.temp_dir, ignore_errors = True) + self.temp_dir = None + + +def boto3_available() -> bool: + """True if boto3 can be imported (without importing it).""" + return find_spec("boto3") is not None + + +def _build_s3_client(s3_config: dict): + """Create a boto3 S3 client from the config dict. + + Uses explicit access keys when provided, otherwise falls back to the + default credential chain (IAM role / instance profile / env / shared creds). + """ + import boto3 # lazy: optional dependency + + region = s3_config.get("region") or "us-east-1" + use_iam_role = bool(s3_config.get("use_iam_role")) + access_key_id = s3_config.get("access_key_id") + secret_access_key = s3_config.get("secret_access_key") + + if not use_iam_role and access_key_id and secret_access_key: + return boto3.client( + "s3", + region_name = region, + aws_access_key_id = access_key_id, + aws_secret_access_key = secret_access_key, + ) + # IAM role / instance profile / ambient credentials + return boto3.client("s3", region_name = region) + + +def _list_dataset_keys(client, bucket: str, prefix: Optional[str]) -> list[str]: + """List object keys under ``prefix`` that have a supported data extension.""" + paginator = client.get_paginator("list_objects_v2") + list_kwargs = {"Bucket": bucket} + if prefix: + list_kwargs["Prefix"] = prefix + + keys: list[str] = [] + for page in paginator.paginate(**list_kwargs): + for obj in page.get("Contents", []): + key = obj["Key"] + if key.endswith("/"): + continue # directory placeholder + if os.path.basename(key).lower() in _IGNORED_METADATA_FILENAMES: + continue + if key.lower().endswith(SUPPORTED_EXTENSIONS): + keys.append(key) + return keys + + +def _extension_family(key: str) -> str: + ext = os.path.splitext(key)[1].lower() + if ext in _JSON_EXTENSIONS: + return "json" + return ext.lstrip(".") + + +def _validate_single_extension_family(keys: list[str]) -> None: + families: list[str] = [] + for key in keys: + family = _extension_family(key) + if family not in families: + families.append(family) + + if len(families) <= 1: + return + + raise ValueError( + "S3 prefix contains mixed dataset formats " + f"({', '.join(families)}). Keep one dataset format under the selected prefix." + ) + + +def _unique_local_path(target_dir: str, filename: str, used_paths: set[str]) -> str: + """Return an unused flattened path for an S3 object basename.""" + stem, ext = os.path.splitext(filename) + candidate = os.path.join(target_dir, filename) + suffix = 1 + while candidate in used_paths or os.path.exists(candidate): + candidate = os.path.join(target_dir, f"{stem}_{suffix}{ext}") + suffix += 1 + used_paths.add(candidate) + return candidate + + +def _raise_if_cancelled(cancel_callback: Optional[Callable[[], bool]]) -> None: + if cancel_callback is not None and cancel_callback(): + raise S3DownloadCancelled("S3 dataset download cancelled") + + +def prepare_s3_dataset_download( + s3_config: dict, + dest_dir: Optional[str] = None, + cancel_callback: Optional[Callable[[], bool]] = None, +) -> S3DatasetDownload: + """Download supported dataset files from S3 to a local directory. + + Returns the local files plus the owned temporary directory, when one was + created. Call ``cleanup()`` after the dataset loader has materialized data. + + Raises ``RuntimeError`` if boto3 is missing, and ``ValueError`` if the + bucket/prefix contains no supported dataset files. + """ + if not boto3_available(): + raise RuntimeError("S3 dataset loading requires boto3. Install it with: pip install boto3") + + bucket = s3_config.get("bucket") + if not bucket: + raise ValueError("s3_config.bucket is required") + prefix = s3_config.get("prefix") + + _raise_if_cancelled(cancel_callback) + client = _build_s3_client(s3_config) + + keys = _list_dataset_keys(client, bucket, prefix) + _raise_if_cancelled(cancel_callback) + if not keys: + where = f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" + raise ValueError( + f"No supported dataset files ({', '.join(SUPPORTED_EXTENSIONS)}) " + f"found under {where}" + ) + + _validate_single_extension_family(keys) + + owns_temp_dir = dest_dir is None + target_dir = dest_dir or tempfile.mkdtemp(prefix = "unsloth_s3_dataset_") + try: + os.makedirs(target_dir, exist_ok = True) + + local_files: list[str] = [] + used_paths: set[str] = set() + for key in keys: + _raise_if_cancelled(cancel_callback) + filename = os.path.basename(key) + local_path = _unique_local_path(target_dir, filename, used_paths) + download_kwargs = {} + if cancel_callback is not None: + download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(cancel_callback) + client.download_file(bucket, key, local_path, **download_kwargs) + _raise_if_cancelled(cancel_callback) + local_files.append(local_path) + except Exception: + if owns_temp_dir: + shutil.rmtree(target_dir, ignore_errors = True) + raise + + logger.info( + "Downloaded %d dataset file(s) from s3://%s/%s to %s", + len(local_files), + bucket, + prefix or "", + target_dir, + ) + return S3DatasetDownload( + files = local_files, + temp_dir = target_dir if owns_temp_dir else None, + ) + + +def download_s3_dataset( + s3_config: dict, + dest_dir: Optional[str] = None, + cancel_callback: Optional[Callable[[], bool]] = None, +) -> list[str]: + download = prepare_s3_dataset_download( + s3_config, + dest_dir = dest_dir, + cancel_callback = cancel_callback, + ) + return download.files diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 57342b2453..018b66ea92 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -2227,6 +2227,7 @@ class UnslothTrainer: dataset_slice_start: int = None, dataset_slice_end: int = None, is_cpt: bool = False, + s3_config: dict = None, ) -> Optional[tuple]: """ Load and prepare a dataset for training. @@ -2237,6 +2238,9 @@ class UnslothTrainer: Returns (dataset_info, eval_dataset) or None on error; eval_dataset may be None if no eval split is available. """ + from core.training.s3_dataset import S3DownloadCancelled + + s3_download = None try: dataset = None eval_dataset = None @@ -2272,6 +2276,22 @@ class UnslothTrainer: return result.dataset + # S3 datasets are downloaded to a local temp dir and then consumed + # through the same local-file path below. + if s3_config and not local_datasets: + from core.training.s3_dataset import prepare_s3_dataset_download + + self._update_progress(status_message = "Downloading dataset from S3...") + s3_download = prepare_s3_dataset_download( + s3_config, + cancel_callback = lambda: self.should_stop, + ) + local_datasets = s3_download.files + if self.should_stop: + logger.info("Stopped during S3 download\n") + return None + logger.info(f"Downloaded {len(local_datasets)} file(s) from S3\n") + if local_datasets: # Use load_dataset() for an Arrow-backed result; in-memory # Dataset.from_list() has no cache and forces num_proc=1 during @@ -2539,10 +2559,16 @@ class UnslothTrainer: return (dataset_info, eval_dataset) + except S3DownloadCancelled: + logger.info("Stopped during S3 download\n") + return None except Exception as e: logger.error(f"Error loading dataset: {e}") self._update_progress(error = str(e)) return None + finally: + if s3_download is not None: + s3_download.cleanup() def _auto_detect_eval_split_from_hf( self, dataset_source: str, subset: str @@ -3367,7 +3393,19 @@ class UnslothTrainer: # ========== PROGRESS TRACKING ========== self.trainer.add_callback(self._create_progress_callback()) - num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + num_samples = None + if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None: + try: + num_samples = len(self.trainer.train_dataset) + except TypeError: + logger.debug( + "train_dataset does not support len(); falling back to " + "raw dataset size for step estimation." + ) + + if num_samples is None: + num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + batch_size = training_args.get("batch_size", 2) total_steps = self._calculate_total_steps( num_samples, @@ -3376,10 +3414,8 @@ class UnslothTrainer: training_args.get("num_epochs", 3), training_args.get("max_steps", 0), ) - self._update_progress(total_steps = total_steps) - # ========== START TRAINING ========== - self._update_progress(status_message = "Starting training...") + self._update_progress(total_steps = total_steps, status_message = "Starting training...") logger.info("Starting training...\n") self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint")) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 2764101d93..6dd42976c7 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -40,6 +40,34 @@ logger = get_logger(__name__) _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") +def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]: + db_config = { + k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"} + } + s3_config = config.get("s3_config") + if hasattr(s3_config, "model_dump"): + s3_config = s3_config.model_dump() + if isinstance(s3_config, dict) and s3_config: + db_config["dataset_source"] = "s3" + db_config["s3_dataset"] = { + "bucket": s3_config.get("bucket"), + "region": s3_config.get("region"), + "prefix": s3_config.get("prefix"), + "use_iam_role": bool(s3_config.get("use_iam_role")), + } + return db_config + + +def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: + if not isinstance(s3_dataset, dict): + return None + bucket = s3_dataset.get("bucket") + if not bucket: + return None + prefix = s3_dataset.get("prefix") + return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" + + def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. @@ -236,6 +264,7 @@ class TrainingBackend: "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), "trust_remote_code": kwargs.get("trust_remote_code", False), "gpu_ids": kwargs.get("gpu_ids"), + "s3_config": kwargs.get("s3_config"), } # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. @@ -309,7 +338,7 @@ class TrainingBackend: self._run_finalized = False self._db_run_created = False self._db_total_steps_set = False - self._db_config = {k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}} + self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() # Assign subprocess handles after state reset. @@ -732,8 +761,11 @@ class TrainingBackend: try: from storage.studio_db import create_run - dataset_name = self._db_config.get("hf_dataset") or next( - iter(self._db_config.get("local_datasets") or []), "unknown" + dataset_name = ( + self._db_config.get("hf_dataset") + or next(iter(self._db_config.get("local_datasets") or []), None) + or _s3_dataset_name(self._db_config.get("s3_dataset")) + or "unknown" ) create_run( id = self.current_job_id, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 18b25cb4fe..1120744a2d 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1336,6 +1336,32 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) + _stop_save = [True] + _stop_requested = [False] + _trainer_ref = [None] + + def _is_stop_requested(): + return _stop_requested[0] + + def _poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 1.0) + if msg and msg.get("type") == "stop": + _stop_save[0] = msg.get("save", True) + _stop_requested[0] = True + trainer = _trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = _poll_stop, daemon = True) + stop_thread.start() + _send("status", status_message = "Loading MLX libraries...") import mlx.core as mx @@ -1515,6 +1541,26 @@ def _run_mlx_training(event_queue, stop_queue, config): elif config.get("local_datasets"): dataset = _load_local(config["local_datasets"]) dataset = _slice(dataset) + elif config.get("s3_config"): + from core.training.s3_dataset import ( + S3DownloadCancelled, + prepare_s3_dataset_download, + ) + + _send("status", status_message = "Downloading dataset from S3...") + try: + s3_download = prepare_s3_dataset_download( + config["s3_config"], + cancel_callback = _is_stop_requested, + ) + try: + dataset = _load_local(s3_download.files) + finally: + s3_download.cleanup() + except S3DownloadCancelled: + _send("complete", output_dir = None, status_message = "Training cancelled") + return + dataset = _slice(dataset) else: raise ValueError("No dataset specified") @@ -1693,6 +1739,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ), ) + _trainer_ref[0] = trainer + if _stop_requested[0]: + trainer.stop_requested = True # Tell the parent eval is configured so the frontend shows the eval chart if eval_dataset is not None and eval_steps_val > 0: @@ -1733,7 +1782,7 @@ def _run_mlx_training(event_queue, stop_queue, config): wandb_token = config.get("wandb_token") if wandb_token: os.environ["WANDB_API_KEY"] = wandb_token - _wandb_sensitive = {"hf_token", "wandb_token"} + _wandb_sensitive = {"hf_token", "wandb_token", "s3_config"} wandb_run = _wandb.init( project = config.get("wandb_project") or "unsloth-mlx", config = {k: v for k, v in config.items() if k not in _wandb_sensitive}, @@ -1835,26 +1884,6 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.add_eval_callback(_on_eval) - # ── 10. Stop signal polling ── - _stop_save = [True] # mutable so thread can update; [save_flag] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - # Safe: pipe permanently broken, no more messages can arrive. - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() - # ── 11. Run training ── gc.collect() mx.synchronize() @@ -2462,7 +2491,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> def _on_progress(progress: TrainingProgress): has_train_loss = progress.step > 0 and progress.loss is not None has_eval_loss = progress.eval_loss is not None - if has_train_loss or has_eval_loss: + if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss: event_queue.put( { "type": "progress", @@ -2546,6 +2575,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> dataset_slice_start = config.get("dataset_slice_start"), dataset_slice_end = config.get("dataset_slice_end"), is_cpt = _is_cpt_for_dataset, + s3_config = config.get("s3_config"), ) if isinstance(dataset_result, tuple): @@ -3010,20 +3040,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> subset = config.get("subset") or None train_split = config.get("train_split", "train") or "train" - if hf_dataset and hf_dataset.strip(): - hf_token = config.get("hf_token", "") - hf_token = hf_token if hf_token and hf_token.strip() else None - dataset = load_dataset( - hf_dataset.strip(), - subset, - split = train_split, - token = hf_token, - ) - elif local_datasets: - # Load local file(s) — mirrors the non-embedding pipeline's directory - # handling so recipe outputs (parquet-files/) work. + def _load_local_embedding_dataset(dataset_paths: list[str]): all_files: list[str] = [] - for dataset_file in local_datasets: + for dataset_file in dataset_paths: file_path = ( dataset_file if os.path.isabs(dataset_file) @@ -3053,17 +3072,58 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> else: all_files.append(file_path) - if all_files: - first_ext = Path(all_files[0]).suffix.lower() - if first_ext in (".json", ".jsonl"): - loader = "json" - elif first_ext == ".csv": - loader = "csv" - elif first_ext == ".parquet": - loader = "parquet" - else: - raise ValueError(f"Unsupported local dataset format: {all_files[0]}") - dataset = load_dataset(loader, data_files = all_files, split = "train") + if not all_files: + raise ValueError("No local dataset files found") + + first_ext = Path(all_files[0]).suffix.lower() + if first_ext in (".json", ".jsonl"): + loader = "json" + elif first_ext == ".csv": + loader = "csv" + elif first_ext == ".parquet": + loader = "parquet" + else: + raise ValueError(f"Unsupported local dataset format: {all_files[0]}") + return load_dataset(loader, data_files = all_files, split = "train") + + if hf_dataset and hf_dataset.strip(): + hf_token = config.get("hf_token", "") + hf_token = hf_token if hf_token and hf_token.strip() else None + dataset = load_dataset( + hf_dataset.strip(), + subset, + split = train_split, + token = hf_token, + ) + elif local_datasets: + dataset = _load_local_embedding_dataset(local_datasets) + elif config.get("s3_config"): + from core.training.s3_dataset import ( + S3DownloadCancelled, + prepare_s3_dataset_download, + ) + + _send_status(event_queue, "Downloading dataset from S3...") + s3_download = None + try: + s3_download = prepare_s3_dataset_download( + config["s3_config"], + cancel_callback = lambda: _should_stop, + ) + dataset = _load_local_embedding_dataset(s3_download.files) + except S3DownloadCancelled: + event_queue.put( + { + "type": "complete", + "output_dir": None, + "status_message": "Training cancelled", + "ts": time.time(), + } + ) + return + finally: + if s3_download is not None: + s3_download.cleanup() else: event_queue.put( { diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index fe3aadfd38..0f7ce6fe34 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -142,13 +142,18 @@ def _compute_all_hf_cache_scans() -> list: logger.warning("Could not scan active HF cache: %s", exc) for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): - extra = extra_fn() - if extra.is_dir() and str(extra.resolve()) not in seen: - seen.add(str(extra.resolve())) - try: - scans.append(scan_cache_dir(cache_dir = str(extra))) - except Exception as exc: - logger.warning("Could not scan HF cache %s: %s", extra, exc) + try: + extra = extra_fn() + # is_dir()/resolve() can raise on an inaccessible path; skip it. + if not extra.is_dir(): + continue + resolved = str(extra.resolve()) + if resolved in seen: + continue + seen.add(resolved) + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc) return scans diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 0d54b435a6..584f82dea0 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -3,14 +3,14 @@ """Pydantic schemas for Export API.""" -from pathlib import Path +from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator from typing import List, Optional, Literal, Dict, Any def _validate_save_directory(value: str) -> str: - """Reject save_directory values that escape the export root.""" + """Validate save_directory — allows absolute paths (user may want a different drive).""" if value is None: raise ValueError("save_directory is required") raw = str(value).strip() @@ -20,15 +20,15 @@ def _validate_save_directory(value: str) -> str: raise ValueError("save_directory may not contain null bytes") if any(ch in raw for ch in ("\r", "\n")): raise ValueError("save_directory may not contain control characters") - if len(raw) > 255: - raise ValueError("save_directory must be <= 255 characters") path = Path(raw).expanduser() - if path.is_absolute(): - raise ValueError( - "save_directory must be a name or relative path under the " - "export root; absolute paths are rejected" - ) - if ".." in path.parts: + path_parts = (*path.parts, *PureWindowsPath(raw).parts, *raw.replace("\\", "/").split("/")) + if any(len(part) > 255 for part in path_parts if part not in ("", ".", "/", "\\")): + raise ValueError("save_directory path components must be <= 255 characters") + if ( + ".." in path.parts + or ".." in PureWindowsPath(raw).parts + or ".." in raw.replace("\\", "/").split("/") + ): raise ValueError("save_directory may not contain '..' segments") return raw diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f94d5366eb..04c63c2247 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -87,6 +87,15 @@ class LoadRequest(BaseModel): "'mtp' or 'mtp+ngram'." ), ) + tensor_parallel: bool = Field( + False, + description = ( + "Split the model across GPUs by tensor (--split-mode tensor) " + "instead of by layer for GGUF models. Only affects multi-GPU " + "setups, where it can make generation significantly faster. " + "No effect on a single GPU. Ignored for non-GGUF models." + ), + ) llama_extra_args: Optional[List[str]] = Field( None, description = ( @@ -160,6 +169,9 @@ class LoadResponse(BaseModel): is_vision: bool = Field(False, description = "Whether model is a vision model") is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether model is a block-diffusion model (DiffusionGemma)" + ) is_audio: bool = Field(False, description = "Whether model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") @@ -171,7 +183,7 @@ class LoadResponse(BaseModel): description = "Whether the model defaults require trust_remote_code to be enabled for loading.", ) context_length: Optional[int] = Field( - None, description = "Model's native context length (from GGUF metadata)" + None, description = "Runtime context length in tokens for the loaded model" ) max_context_length: Optional[int] = Field( None, description = "Maximum context length currently available on this hardware" @@ -224,6 +236,10 @@ class LoadResponse(BaseModel): "None when the platform default is in effect." ), ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) class UnloadResponse(BaseModel): @@ -273,6 +289,9 @@ class InferenceStatusResponse(BaseModel): ) is_vision: bool = Field(False, description = "Whether the active model is a vision model") is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)" + ) gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)") is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") @@ -339,6 +358,10 @@ class InferenceStatusResponse(BaseModel): "None when the platform default is in effect." ), ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( @@ -690,6 +713,10 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.", ) + confirm_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.", + ) auto_heal_tool_calls: Optional[bool] = Field( True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", @@ -926,6 +953,12 @@ class ChatCompletionRequest(BaseModel): return self +class ToolConfirmRequest(BaseModel): + session_id: Optional[str] = None + approval_id: Optional[str] = None + decision: Literal["allow", "deny"] = "deny" + + # ── OpenAI shell-tool container management ───────────────────── @@ -1304,6 +1337,23 @@ class ResponsesOutputMessage(BaseModel): content: list[ResponsesOutputTextContent] = Field(default_factory = list) +class ResponsesOutputReasoningContent(BaseModel): + """A reasoning text content block inside a reasoning output item.""" + + type: Literal["reasoning_text"] = "reasoning_text" + text: str + + +class ResponsesOutputReasoning(BaseModel): + """A top-level reasoning output item in the Responses API response.""" + + type: Literal["reasoning"] = "reasoning" + id: str = Field(default_factory = lambda: f"rs_{uuid.uuid4().hex[:12]}") + status: Literal["completed", "in_progress", "incomplete"] = "completed" + summary: list = Field(default_factory = list) + content: Optional[list[ResponsesOutputReasoningContent]] = None + + class ResponsesOutputFunctionCall(BaseModel): """A function-call output item in the Responses API response. @@ -1318,7 +1368,11 @@ class ResponsesOutputFunctionCall(BaseModel): status: Literal["completed", "in_progress", "incomplete"] = "completed" -ResponsesOutputItem = Union[ResponsesOutputMessage, ResponsesOutputFunctionCall] +ResponsesOutputItem = Union[ + ResponsesOutputMessage, + ResponsesOutputReasoning, + ResponsesOutputFunctionCall, +] class ResponsesUsage(BaseModel): diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 8f90e19df6..5a75246c07 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -108,6 +108,9 @@ class ProviderTestRequest(BaseModel): base_url: Optional[str] = Field( None, description = "Custom base URL (overrides registry default)" ) + model_id: Optional[str] = Field( + None, description = "Model ID for providers that need a chat probe" + ) class ProviderTestResult(BaseModel): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 4a3303c734..ca04591178 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -29,6 +29,43 @@ _MIN_VISION_IMAGE_SIZE = 256 _MAX_VISION_IMAGE_SIZE = 2048 +class S3Config(BaseModel): + """S3 bucket configuration for loading datasets from AWS S3""" + + # Accept both snake_case and the frontend's camelCase field names. + model_config = ConfigDict(populate_by_name = True) + + bucket: str = Field(..., description = "S3 bucket name") + region: str = Field("us-east-1", description = "AWS region") + prefix: Optional[str] = Field(None, description = "Optional path prefix within bucket") + access_key_id: Optional[str] = Field( + None, + alias = "accessKeyId", + description = "AWS access key ID (optional if using IAM role)", + ) + secret_access_key: Optional[str] = Field( + None, + alias = "secretAccessKey", + description = "AWS secret access key (optional if using IAM role)", + ) + use_iam_role: bool = Field( + False, + alias = "useIamRole", + description = "Use IAM role credentials instead of access keys", + ) + + @model_validator(mode = "after") + def _check_credentials(self) -> "S3Config": + # Require either IAM role auth or a full key pair so credentials are + # never half-configured. + if not self.use_iam_role and not (self.access_key_id and self.secret_access_key): + raise ValueError( + "s3_config requires either use_iam_role=True or both " + "access_key_id and secret_access_key" + ) + return self + + def _parse_lr(v: Any) -> float: """Parse learning_rate as a positive float strictly below _MAX_LR_VALUE.""" if v is None: @@ -338,6 +375,12 @@ class TrainingStartRequest(BaseModel): description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", ) + # S3 dataset source configuration + s3_config: Optional[S3Config] = Field( + None, + description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", + ) + @model_validator(mode = "after") def _check_steps_or_epochs(self) -> "TrainingStartRequest": # Each accepts 0 as "use the other"; both 0 means nothing to train. diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 85294114b1..6efe91d448 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -68,3 +68,9 @@ trl>=0.18.2,!=0.19.0,<=0.24.0 sentence-transformers cut_cross_entropy pillow + +# RAG store + document parsing, mirroring studio.txt. Pinned here because +# this file installs --no-deps; without them Studio runs with RAG disabled. +sqlite-vec==0.1.9 +pymupdf==1.27.2.3 +python-docx==1.2.0 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 9bffb81549..96fef60471 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -17,6 +17,7 @@ structlog>=24.1.0 diceware ddgs cryptography>=42.0.0 +boto3>=1.34.0 # optional: S3 dataset loading httpx>=0.27.0 fastmcp>=3.0.2 # RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index ef37daa158..f7b3a56a71 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -154,19 +154,41 @@ async def get_export_status(current_subject: str = Depends(get_current_subject)) ) +def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]: + """Best-effort registration so absolute exports show up in local scans.""" + try: + from storage.studio_db import add_scan_folder + folder = add_scan_folder(str(path)) + return True, str(folder.get("path") or path) + except Exception as exc: + logger.warning("Could not register export scan folder %s: %s", path, exc) + return False, None + + def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: - """Return the export path relative to exports_root, hiding the install path.""" + """Return relative export paths, keeping external absolute paths visible.""" if not output_path: return None try: from utils.paths.storage_roots import exports_root + path = Path(output_path) + # If it's outside exports_root, return the full absolute path + # so users can find their files on a different drive. + if path.is_absolute(): + try: + path.resolve().relative_to(exports_root().resolve()) + except ValueError: + registered, registered_path = _try_register_external_export(path) + return { + "output_path": str(path), + "scan_folder_registered": registered, + "scan_folder_path": registered_path, + } rel = os.path.relpath(output_path, exports_root()) - if rel.startswith(".."): - rel = os.path.basename(output_path) return {"output_path": rel} except Exception: - return {"output_path": os.path.basename(output_path)} + return {"output_path": output_path} @router.post("/export/merged", response_model = ExportOperationResponse) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3333db9d07..fa708d0e39 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -29,6 +29,16 @@ from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +def _positive_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int > 0 else None + + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -586,9 +596,11 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, ) + from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( @@ -617,9 +629,11 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, ) + from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( @@ -645,6 +659,7 @@ from models.inference import ( ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, + ToolConfirmRequest, ChatMessage, ChunkChoice, ChoiceDelta, @@ -667,6 +682,8 @@ from models.inference import ( ResponsesFunctionCallOutputInputItem, ResponsesOutputTextContent, ResponsesOutputMessage, + ResponsesOutputReasoning, + ResponsesOutputReasoningContent, ResponsesOutputFunctionCall, ResponsesUsage, ResponsesResponse, @@ -692,10 +709,12 @@ from core.inference.anthropic_compat import ( AnthropicPassthroughEmitter, ) from auth.authentication import get_current_subject +from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key from core.inference.providers import get_provider_info, get_base_url from core.inference.external_provider import ExternalProviderClient +from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db from utils.utils import safe_error_detail, log_and_http_error @@ -1147,9 +1166,34 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: return value -def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaCppBackend) -> bool: +def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: + """Whether an inherited --split-mode should be stripped on reload. + + The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ + layer modes, so only strip when the toggle overrides it: tensor being turned + on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes + survive. Shared by the inheritance strip and the already-loaded stale check + so they agree on what reload would do. + """ + fields_set = getattr(request, "model_fields_set", set()) + return "tensor_parallel" in fields_set and ( + request.tensor_parallel or resolve_tensor_parallel(backend_extra, False) + ) + + +def _request_matches_loaded_settings( + request: LoadRequest, + llama_backend: LlamaCppBackend, + effective_chat_template_override: Optional[str] = None, +) -> bool: """True iff every runtime setting on the request matches the loaded server. - Caller has already checked model+variant+is_loaded. See #5401.""" + Caller has already checked model+variant+is_loaded. See #5401. + + ``effective_chat_template_override`` is the resolved template that will be + launched (user override, else a bundled family template such as the + gemma-4 override), so the dedup compares against what the backend actually + holds rather than the raw request field. Defaults to the request field for + callers that do not resolve a bundled override.""" # Compare requested n_ctx (not effective) so VRAM-cap doesn't mask an # Auto-vs-explicit slider flip. if request.max_seq_length != llama_backend.requested_n_ctx: @@ -1158,6 +1202,24 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC llama_backend.cache_type_kv ): return False + # Reconcile a user --split-mode in extras into the effective tensor state. + # When the request omits llama_extra_args ("inherit"), compare using the + # stored extras stripped the way the reload strips them, so an extras-driven + # tensor load isn't seen as a mismatch that needlessly reloads the server. + backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + effective_extra = ( + request.llama_extra_args + if request.llama_extra_args is not None + else strip_shadowing_flags( + backend_extra, + strip_split_mode = _should_strip_split_mode(request, backend_extra), + ) + ) + if ( + resolve_tensor_parallel(effective_extra, request.tensor_parallel) + != llama_backend.tensor_parallel + ): + return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare # the real requested mode -- coercing vision to ``off`` here used to @@ -1171,15 +1233,29 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None: if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0): return False - if (request.chat_template_override or None) != (llama_backend.chat_template_override or None): + _effective_cto = ( + effective_chat_template_override + if effective_chat_template_override is not None + else request.chat_template_override + ) + if (_effective_cto or None) != (llama_backend.chat_template_override or None): return False # llama_extra_args=None means "inherit"; only an explicit differing list # forces a reload. On the inherit path, refuse to match if stored extras # contain any shadow flag, so the reload path strips them rather than - # leaving a stale override in effect. - backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra: + # Mirror the reload's conditional split-mode strip, so a preserved + # non-tensor mode (row/none/layer) isn't seen as stale and doesn't + # trigger a needless reload of a healthy server. + if ( + backend_extra + and strip_shadowing_flags( + backend_extra, + strip_split_mode = _should_strip_split_mode(request, backend_extra), + ) + != backend_extra + ): return False else: if list(request.llama_extra_args) != backend_extra: @@ -1288,6 +1364,17 @@ async def load_model( # Version switching is handled by the subprocess-based inference # backend -- no ensure_transformers_version() needed here. + # Resolve the effective chat-template override once, up front: an + # explicit user override, else a bundled family template (e.g. the + # gemma-4 override that ships preserve_thinking without re-downloading + # quants), else None. Used for both the reload-dedup check below and the + # load_model calls, so the live backend state and the incoming request + # compare against the same template text. + effective_chat_template_override = resolve_effective_chat_template_override( + model_identifier = model_identifier, + user_override = request.chat_template_override, + ) + # ── Already-loaded check: skip reload if the exact model is active ── backend = get_inference_backend() llama_backend = get_llama_cpp_backend() @@ -1305,7 +1392,9 @@ async def load_model( and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() # Match runtime settings so Apply isn't dropped (#5401). - and _request_matches_loaded_settings(request, llama_backend) + and _request_matches_loaded_settings( + request, llama_backend, effective_chat_template_override + ) # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) ): @@ -1330,6 +1419,7 @@ async def load_model( is_vision = llama_backend._is_vision, is_lora = False, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, has_audio_input = getattr(llama_backend, "_has_audio_input", False), @@ -1348,6 +1438,7 @@ async def load_model( chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, ) else: if ( @@ -1390,6 +1481,7 @@ async def load_model( reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], + context_length = _positive_int_or_none(_model_info.get("context_length")), chat_template = _chat_template, ) @@ -1438,18 +1530,23 @@ async def load_model( # parse against a freshly-supplied first-class field. if request.llama_extra_args is None and llama_backend.extra_args: source = llama_backend.extra_args_source - # Compare against the resolved variant, not the request field: - # callers commonly omit gguf_variant for local ``.gguf`` paths - # and HF auto-pick flows. ``config.gguf_variant`` is the variant - # load_model was actually invoked with (see HF / local branches - # below), so both sides key off the same string. - resolved_variant = config.gguf_variant - same_source = bool( - source - and source[0] - and source[0].lower() == model_identifier.lower() - and (source[1] or "").lower() == (resolved_variant or "").lower() + # Compare against the resolved variant, not the request + # field: callers commonly omit gguf_variant for local + # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ + # variant`` is the variant load_model was actually + # invoked with (see the HF / local branches below), so + # both sides of the comparison key off the same string. + resolved_variant = (config.gguf_variant or "").lower() + request_variant = (request.gguf_variant or "").lower() + stored_variant = (source[1] or "").lower() if source else "" + same_model = bool( + source and source[0] and source[0].lower() == model_identifier.lower() ) + if request.gguf_variant: + variant_mismatch = request_variant != stored_variant + else: + variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + same_source = same_model and not variant_mismatch if not same_source: logger.info( "Not inheriting llama_extra_args: stored args came from %s, loading %s", @@ -1462,7 +1559,13 @@ async def load_model( else: # Strip only the groups whose first-class field was set by # the caller, so an inherited --chat-template-file survives - # an Apply that omits chat_template_override. + # an Apply that omits chat_template_override. A bundled family + # template (e.g. the gemma-4 override) is an effective + # first-class template setting even when the raw request + # omits chat_template_override, so strip the inherited + # --chat-template-file in that case too -- otherwise the stale + # extra arg (appended last) shadows the bundled template while + # Studio reports the bundled template's capabilities. fields_set = getattr(request, "model_fields_set", set()) stripped = strip_shadowing_flags( llama_backend.extra_args, @@ -1471,7 +1574,13 @@ async def load_model( strip_spec = ( "speculative_type" in fields_set or "spec_draft_n_max" in fields_set ), - strip_template = "chat_template_override" in fields_set, + strip_template = ( + "chat_template_override" in fields_set + or effective_chat_template_override is not None + ), + strip_split_mode = _should_strip_split_mode( + request, llama_backend.extra_args + ), ) try: extra_llama_args = validate_extra_args(stripped) @@ -1479,8 +1588,7 @@ async def load_model( # Shouldn't happen on already-validated args; degrade to # no-extras rather than 400 if managed flags changed. logger.warning( - "Stored llama_extra_args failed revalidation; " - "loading without them: %s", + "Stored llama_extra_args failed revalidation; loading without them: %s", stripped, ) extra_llama_args = [] @@ -1497,22 +1605,25 @@ async def load_model( # during the (potentially long) GGUF download + llama-server start. _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + # Load kwargs common to HF and local modes; the two differ only by + # the model-source args (hf_repo/-token vs gguf_path/mmproj). + _common_load_kwargs = dict( + model_identifier = config.identifier, + is_vision = config.is_vision, + n_ctx = request.max_seq_length, + chat_template_override = effective_chat_template_override, + cache_type_kv = request.cache_type_kv, + speculative_type = request.speculative_type, + spec_draft_n_max = request.spec_draft_n_max, + n_parallel = _n_parallel, + extra_args = extra_llama_args, + ) if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server - success = await asyncio.to_thread( - llama_backend.load_model, + _source_load_kwargs = dict( hf_repo = config.gguf_hf_repo, hf_variant = config.gguf_variant, hf_token = request.hf_token, - model_identifier = config.identifier, - is_vision = config.is_vision, - n_ctx = request.max_seq_length, - chat_template_override = request.chat_template_override, - cache_type_kv = request.cache_type_kv, - speculative_type = request.speculative_type, - spec_draft_n_max = request.spec_draft_n_max, - n_parallel = _n_parallel, - extra_args = extra_llama_args, ) else: # Local mode: llama-server loads via -m @@ -1531,8 +1642,7 @@ async def load_model( except HTTPException as exc: logger.warning("Dropping MTP drafter for native load: %s", exc.detail) config.gguf_mtp_file = None - success = await asyncio.to_thread( - llama_backend.load_model, + _source_load_kwargs = dict( gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, mtp_draft_path = config.gguf_mtp_file, @@ -1540,17 +1650,36 @@ async def load_model( # the same string the inheritance check at the top of /load # uses (#5401 followup). hf_variant = config.gguf_variant, - model_identifier = config.identifier, - is_vision = config.is_vision, - n_ctx = request.max_seq_length, - chat_template_override = request.chat_template_override, - cache_type_kv = request.cache_type_kv, - speculative_type = request.speculative_type, - spec_draft_n_max = request.spec_draft_n_max, - n_parallel = _n_parallel, - extra_args = extra_llama_args, ) + # Run a single load attempt with the given tensor flag + extras. + async def _attempt_gguf_load( + tensor_parallel: bool, attempt_extra_args: Optional[list[str]] + ) -> bool: + attempt_kwargs = { + **_common_load_kwargs, + "extra_args": attempt_extra_args, + } + return await asyncio.to_thread( + llama_backend.load_model, + **_source_load_kwargs, + **attempt_kwargs, + tensor_parallel = tensor_parallel, + ) + + # Tensor parallelism is arch-gated in llama.cpp and crashes some loads + # outright (e.g. Gemma 3n aborts with a GGML_ASSERT). The helper auto- + # falls back to layer split so the checkbox never blocks a model from + # loading; the response reports the backend's actual tensor_parallel + # state so the UI toggle reflects the fallback. + success = await load_with_tensor_fallback( + _attempt_gguf_load, + requested_tensor = request.tensor_parallel, + extra_args = extra_llama_args, + label = config.identifier, + cancelled = llama_backend.load_cancelled, + ) + if not success: raise HTTPException( status_code = 500, @@ -1578,6 +1707,7 @@ async def load_model( is_vision = llama_backend.is_vision, is_lora = False, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, has_audio_input = llama_backend._has_audio_input, @@ -1595,6 +1725,7 @@ async def load_model( chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -1731,6 +1862,7 @@ async def load_model( reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], + context_length = _positive_int_or_none(_model_info.get("context_length")), chat_template = _chat_template, ) @@ -1921,6 +2053,20 @@ async def cancel_inference(request: Request, current_subject: str = Depends(get_ return {"cancelled": n} +@studio_router.post("/tool-confirm") +async def confirm_tool_call( + request: ToolConfirmRequest, current_subject: str = Depends(get_current_subject) +): + matched = resolve_tool_decision( + request.approval_id, + request.decision, + session_id = request.session_id, + ) + if not matched: + raise HTTPException(status_code = 404, detail = "No pending tool call confirmation") + return {"resolved": True} + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, current_subject: str = Depends(get_current_subject) @@ -2040,11 +2186,27 @@ async def get_status(current_subject: str = Depends(get_current_subject)): _display_model_id = os.path.basename(_model_id) _inference_cfg = load_inference_config(_model_id) if _model_id else None _audio_type = getattr(llama_backend, "_audio_type", None) + # Don't surface Studio's auto-applied bundled family template (e.g. the + # gemma-4 override) as a user-authored override: the frontend adopts + # status.chat_template_override as editable state and would otherwise + # re-send it as an explicit override for a later, unrelated model. Only + # expose a genuine user override. + _reported_chat_template_override = llama_backend.chat_template_override + _auto_chat_template_override = resolve_effective_chat_template_override( + model_identifier = _model_id, + user_override = None, + ) + if ( + _auto_chat_template_override is not None + and _reported_chat_template_override == _auto_chat_template_override + ): + _reported_chat_template_override = None return InferenceStatusResponse( active_model = _display_model_id, model_identifier = None if _native_grant_backed else _model_id, is_vision = llama_backend.is_vision, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, gguf_variant = llama_backend.hf_variant, is_audio = getattr(llama_backend, "_is_audio", False), audio_type = _audio_type, @@ -2065,9 +2227,10 @@ async def get_status(current_subject: str = Depends(get_current_subject)): max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, cache_type_kv = llama_backend.cache_type_kv, - chat_template_override = llama_backend.chat_template_override, + chat_template_override = _reported_chat_template_override, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, @@ -2119,6 +2282,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): reasoning_always_on = _sf_flags["reasoning_always_on"], supports_preserve_thinking = _sf_flags["supports_preserve_thinking"], supports_tools = _sf_flags["supports_tools"], + context_length = _positive_int_or_none(model_info.get("context_length")), chat_template = chat_template, llama_cpp_supports_mtp = _supports_mtp, llama_cpp_prebuilt_stale = _stale, @@ -3178,6 +3342,22 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + if payload.confirm_tool_calls and ( + payload.enable_tools is True + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls is only supported for local streaming tools.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) if _wants_multiple_choices(payload): _raise_unsupported_n("external provider chat completions") return await _proxy_to_external_provider(payload, request) @@ -3549,6 +3729,16 @@ async def openai_chat_completions( use_tools = False if use_tools: + if payload.confirm_tool_calls and not payload.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── @@ -3619,6 +3809,7 @@ async def openai_chat_completions( session_id = payload.session_id, rag_scope = payload.rag_scope, disable_parallel_tool_use = payload.parallel_tool_calls is False, + confirm_tool_calls = bool(payload.confirm_tool_calls), ) _tool_sentinel = object() @@ -3628,6 +3819,7 @@ async def openai_chat_completions( _tracker.__enter__() async def gguf_tool_stream(): + gen = None try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -3750,6 +3942,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + if gen is not None: + try: + gen.close() + except (RuntimeError, ValueError): + pass _tracker.__exit__(None, None, None) return StreamingResponse( @@ -3833,6 +4030,10 @@ async def openai_chat_completions( _stream_usage = cumulative.get("usage") _stream_timings = cumulative.get("timings") _stream_finish = cumulative.get("finish_reason") + elif cumulative.get("type") == "diffusion_frame": + # Diffusion frame (per-step canvas): pass through as a raw SSE line on the + # tool_status channel. No assistant text, so it never enters the cumulative diff. + yield f"data: {json.dumps(cumulative)}\n\n" else: logger.warning( "gguf_stream_chunks: unexpected dict event: %s", @@ -4073,6 +4274,16 @@ async def openai_chat_completions( _sf_use_tools = False if _sf_use_tools: + if payload.confirm_tool_calls and not payload.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) _sf_nudge = _build_tool_action_nudge( tools = _sf_tools_to_use, model_name = model_name, @@ -4149,6 +4360,7 @@ async def openai_chat_completions( else 300, session_id = payload.session_id, rag_scope = payload.rag_scope, + confirm_tool_calls = bool(payload.confirm_tool_calls), use_adapter = payload.use_adapter, stats_holder = _sf_stats_holder, ) @@ -4159,6 +4371,7 @@ async def openai_chat_completions( _sf_tracker.__enter__() async def sf_tool_stream(): + gen = None try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -4275,6 +4488,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + if gen is not None: + try: + gen.close() + except (RuntimeError, ValueError): + pass _sf_tracker.__exit__(None, None, None) if payload.stream: @@ -4625,28 +4843,38 @@ def _openai_model_objects() -> list[dict]: "created": _created, "owned_by": "local", } - # Extension fields: the real per-request window (post /props readback) - # so clients can budget/compact against the enforced limit. - if llama_backend.context_length: - entry["context_length"] = llama_backend.context_length - if llama_backend.max_context_length: - entry["max_context_length"] = llama_backend.max_context_length + _ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None)) + if _ctx is not None: + entry["context_length"] = _ctx + _max_ctx = _positive_int_or_none(getattr(llama_backend, "max_context_length", None)) + if _max_ctx is not None: + entry["max_context_length"] = _max_ctx + _native_ctx = _positive_int_or_none(getattr(llama_backend, "native_context_length", None)) + if _native_ctx is not None: + entry["native_context_length"] = _native_ctx models.append(entry) # Check Unsloth backend backend = get_inference_backend() if backend.active_model_name: + model_info = backend.models.get(backend.active_model_name, {}) entry = { "id": backend.active_model_name, "object": "model", "created": _created, "owned_by": "local", } - _sf_ctx = getattr(backend, "context_length", None) or getattr( - backend, "max_seq_length", None - ) - if _sf_ctx: - entry["context_length"] = _sf_ctx + _ctx = _positive_int_or_none(model_info.get("context_length")) + if _ctx is None: + for _candidate in ( + getattr(backend, "context_length", None), + getattr(backend, "max_seq_length", None), + ): + _ctx = _positive_int_or_none(_candidate) + if _ctx is not None: + break + if _ctx is not None: + entry["context_length"] = _ctx models.append(entry) return models @@ -4908,6 +5136,149 @@ def _responses_tool_output_text(output: Union[str, list]) -> str: return "(no output)" +_RESPONSES_THINK_OPEN = "" +_RESPONSES_THINK_CLOSE = "" +_RESPONSES_REASONING_EFFORTS = {"none", "minimal", "low", "medium", "high", "max", "xhigh"} + + +def _coerce_responses_reasoning_text(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, list): + return "".join(_coerce_responses_reasoning_text(part) for part in value) + if isinstance(value, dict): + for key in ("text", "reasoning_text", "content"): + text = _coerce_responses_reasoning_text(value.get(key)) + if text: + return text + return "" + return json.dumps(value) + + +def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int: + """Number of trailing chars to retain because they may start a marker.""" + for size in range(min(len(text), max(len(m) for m in markers) - 1), 0, -1): + suffix = text[-size:] + if any(marker.startswith(suffix) for marker in markers): + return size + return 0 + + +class _ResponsesReasoningExtractor: + """Split local markup into Responses reasoning and visible text.""" + + def __init__(self, *, parse_think_markers: bool = False) -> None: + self._buffer = "" + self._in_reasoning = False + self._parse_think_markers = parse_think_markers + + def feed( + self, + text: str = "", + reasoning_content: Any = None, + ) -> tuple[str, str]: + reasoning_parts: list[str] = [] + visible_parts: list[str] = [] + structured_reasoning = _coerce_responses_reasoning_text(reasoning_content) + if structured_reasoning: + reasoning_parts.append(structured_reasoning) + if text: + self._buffer += text + if not self._parse_think_markers: + visible_parts.append(self._buffer) + self._buffer = "" + return "".join(reasoning_parts), "".join(visible_parts) + + while self._buffer: + if self._in_reasoning: + close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) + if close_idx != -1: + reasoning_parts.append(self._buffer[:close_idx]) + self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] + self._in_reasoning = False + continue + keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + if keep == len(self._buffer): + break + reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) + self._buffer = self._buffer[-keep:] if keep else "" + break + + open_idx = self._buffer.find(_RESPONSES_THINK_OPEN) + close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) + if close_idx != -1 and (open_idx == -1 or close_idx < open_idx): + visible_parts.append(self._buffer[:close_idx]) + self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] + continue + if open_idx != -1: + visible_parts.append(self._buffer[:open_idx]) + self._buffer = self._buffer[open_idx + len(_RESPONSES_THINK_OPEN) :] + self._in_reasoning = True + continue + + keep = _responses_marker_holdback( + self._buffer, + (_RESPONSES_THINK_OPEN, _RESPONSES_THINK_CLOSE), + ) + if keep == len(self._buffer): + break + visible_parts.append(self._buffer[:-keep] if keep else self._buffer) + self._buffer = self._buffer[-keep:] if keep else "" + break + + return "".join(reasoning_parts), "".join(visible_parts) + + def finish(self) -> tuple[str, str]: + if not self._buffer: + return "", "" + remaining = self._buffer + self._buffer = "" + if not self._parse_think_markers: + return "", remaining + if self._in_reasoning: + self._in_reasoning = False + return remaining, "" + return "", remaining.replace(_RESPONSES_THINK_CLOSE, "") + + +def _extract_responses_reasoning( + text: str = "", + reasoning_content: Any = None, + *, + parse_think_markers: bool = False, +) -> tuple[str, str]: + extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers) + reasoning, visible = extractor.feed(text, reasoning_content) + final_reasoning, final_visible = extractor.finish() + return reasoning + final_reasoning, visible + final_visible + + +def _responses_should_parse_think_markers( + chat_req: ChatCompletionRequest, llama_backend: Any = None +) -> bool: + if llama_backend is not None and getattr(llama_backend, "is_loaded", False): + if getattr(llama_backend, "reasoning_always_on", False): + return True + if not getattr(llama_backend, "supports_reasoning", False): + return False + if chat_req.enable_thinking is True: + return True + return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none") + + +def _responses_reasoning_output_item(reasoning_text: str, item_id: Optional[str] = None) -> dict: + kwargs: dict[str, Any] = { + "status": "completed", + "summary": [], + "content": [ResponsesOutputReasoningContent(text = reasoning_text)], + } + if item_id is not None: + kwargs["id"] = item_id + return ResponsesOutputReasoning(**kwargs).model_dump() + + def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: """Convert a ResponsesRequest's ``input`` into a Chat-format ``ChatMessage`` list. @@ -5066,6 +5437,33 @@ def _build_chat_request( if payload.parallel_tool_calls is not None: chat_kwargs["parallel_tool_calls"] = payload.parallel_tool_calls + # ``chat_template_kwargs`` (e.g. ``{"enable_thinking": true}``) arrives via + # the Responses extra-body: ResponsesRequest has ``extra="allow"``, so the + # OpenAI SDK's ``extra_body`` spread lands the dict in ``model_extra``. The + # downstream Chat Completions paths consume the typed ``enable_thinking`` + # field -- the non-streaming path lifts it in ``openai_chat_completions`` + # only when it is still ``None``, and the streaming pass-through reads + # ``payload.enable_thinking`` directly -- so lift it here, mirroring that + # handler, to cover both Responses paths. + explicit_enable_thinking = False + _extra = getattr(payload, "model_extra", None) + if isinstance(_extra, dict): + _tpl_kw = _extra.get("chat_template_kwargs") + if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: + chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"]) + explicit_enable_thinking = True + + if isinstance(payload.reasoning, dict): + effort = payload.reasoning.get("effort") + if isinstance(effort, str) and effort in _RESPONSES_REASONING_EFFORTS: + if not explicit_enable_thinking: + chat_kwargs["reasoning_effort"] = effort + chat_kwargs["enable_thinking"] = effort != "none" + elif chat_kwargs.get("enable_thinking") is False: + chat_kwargs["reasoning_effort"] = "none" + elif effort != "none": + chat_kwargs["reasoning_effort"] = effort + return ChatCompletionRequest(**chat_kwargs) @@ -5109,10 +5507,18 @@ async def _responses_non_streaming( choices = body.get("choices", []) text = "" + reasoning_text = "" tool_calls: list[dict] = [] if choices: msg = choices[0].get("message", {}) or {} - text = msg.get("content", "") or "" + raw_content = msg.get("content", "") or "" + raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) + llama_backend = get_llama_cpp_backend() + reasoning_text, text = _extract_responses_reasoning( + raw_text, + msg.get("reasoning_content"), + parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), + ) tool_calls = msg.get("tool_calls") or [] usage_data = body.get("usage", {}) @@ -5126,6 +5532,10 @@ async def _responses_non_streaming( # the model produced content, so clients expecting a pure tool-call turn # (finish_reason="tool_calls") don't see a spurious empty message item. output_items: list[dict] = [] + if reasoning_text and not text and not tool_calls: + text = reasoning_text + if reasoning_text: + output_items.append(_responses_reasoning_output_item(reasoning_text)) if text: msg_id = f"msg_{uuid.uuid4().hex[:12]}" output_items.append( @@ -5173,16 +5583,15 @@ async def _responses_stream( avoids that. Non-GGUF falls back to the wrapper (which doesn't use httpx, so the issue doesn't apply). - Text deltas arrive as ``response.output_text.delta`` on a single - ``message`` output item at ``output_index=0``. Each tool call from + Output items are allocated as upstream deltas appear. Reasoning/text deltas + open top-level ``reasoning`` / ``message`` items; each tool call from ``delta.tool_calls[]`` is promoted to its own top-level ``function_call`` - output item (one per distinct ``tool_calls[].index``) and relayed as + item (one per distinct ``tool_calls[].index``) and relayed as ``response.function_call_arguments.delta`` / ``.done`` events so clients (Codex, OpenAI Python SDK) can reconstruct the call incrementally and reply with a ``function_call_output`` item next turn. """ resp_id = f"resp_{uuid.uuid4().hex[:12]}" - msg_id = f"msg_{uuid.uuid4().hex[:12]}" created_at = int(time.time()) chat_req = _build_chat_request(payload, messages, stream = True) @@ -5222,61 +5631,166 @@ async def _responses_stream( async def event_generator(): full_text = "" + full_reasoning = "" input_tokens = 0 output_tokens = 0 + extractor = _ResponsesReasoningExtractor( + parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) + ) + reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} - # Text message lives at output_index 0; tool calls claim 1, 2, ... - next_output_index = 1 + next_output_index = 0 + + def _sse(event_name: str, payload: dict) -> str: + return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n" + + def _claim_output_index() -> int: + nonlocal next_output_index + output_index = next_output_index + next_output_index += 1 + return output_index + + def _ensure_reasoning_open() -> list[str]: + if reasoning_state["opened"]: + return [] + reasoning_state["output_index"] = _claim_output_index() + reasoning_state["item_id"] = f"rs_{uuid.uuid4().hex[:12]}" + reasoning_state["opened"] = True + output_index = reasoning_state["output_index"] + item_id = reasoning_state["item_id"] + return [ + _sse( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "reasoning", + "id": item_id, + "status": "in_progress", + "summary": [], + "content": [], + }, + }, + ), + _sse( + "response.content_part.added", + { + "type": "response.content_part.added", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "part": {"type": "reasoning_text", "text": ""}, + }, + ), + ] + + def _ensure_message_open() -> list[str]: + if message_state["opened"]: + return [] + message_state["output_index"] = _claim_output_index() + message_state["item_id"] = f"msg_{uuid.uuid4().hex[:12]}" + message_state["opened"] = True + output_index = message_state["output_index"] + item_id = message_state["item_id"] + return [ + _sse( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "message", + "id": item_id, + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + ), + _sse( + "response.content_part.added", + { + "type": "response.content_part.added", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + }, + ), + ] def _snapshot_output() -> list[dict]: """Snapshot of all completed output items for response.completed.""" - items: list[dict] = [ - { - "type": "message", - "id": msg_id, - "status": "completed", - "role": "assistant", - "content": [ + indexed_items: list[tuple[int, dict]] = [] + if reasoning_state["opened"]: + indexed_items.append( + ( + reasoning_state["output_index"], { - "type": "output_text", - "text": full_text, - "annotations": [], - } - ], - } - ] - for st in sorted(tool_call_state.values(), key = lambda s: s["output_index"]): - items.append( - { - "type": "function_call", - "id": st["item_id"], - "status": "completed", - "call_id": st["call_id"], - "name": st["name"], - "arguments": st["arguments"], - } + "type": "reasoning", + "id": reasoning_state["item_id"], + "status": "completed", + "summary": [], + "content": [{"type": "reasoning_text", "text": full_reasoning}], + }, + ) ) - return items + if message_state["opened"]: + indexed_items.append( + ( + message_state["output_index"], + { + "type": "message", + "id": message_state["item_id"], + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": full_text, + "annotations": [], + } + ], + }, + ) + ) + for st in tool_call_state.values(): + indexed_items.append( + ( + st["output_index"], + { + "type": "function_call", + "id": st["item_id"], + "status": "completed", + "call_id": st["call_id"], + "name": st["name"], + "arguments": st["arguments"], + }, + ) + ) + return [item for _, item in sorted(indexed_items, key = lambda pair: pair[0])] # ── Preamble events ── - yield f"event: response.created\ndata: {json.dumps({'type': 'response.created', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'in_progress', 'model': payload.model, 'output': [], 'usage': {'input_tokens': 0, 'output_tokens': 0, 'total_tokens': 0}}})}\n\n" - - # output_item.added (text message at output_index 0) - output_item = { - "type": "message", - "id": msg_id, - "status": "in_progress", - "role": "assistant", - "content": [], - } - yield f"event: response.output_item.added\ndata: {json.dumps({'type': 'response.output_item.added', 'output_index': 0, 'item': output_item})}\n\n" - - # content_part.added - content_part = {"type": "output_text", "text": "", "annotations": []} - yield f"event: response.content_part.added\ndata: {json.dumps({'type': 'response.content_part.added', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': content_part})}\n\n" + yield _sse( + "response.created", + { + "type": "response.created", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "in_progress", + "model": payload.model, + "output": [], + "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + }, + }, + ) # ── Direct httpx lifecycle to llama-server ── # Full same-task open + close, same pattern as @@ -5293,7 +5807,21 @@ async def _responses_stream( resp = await client.send(req, stream = True) except httpx.RequestError as e: logger.error("responses stream: upstream unreachable: %s", e) - yield f"event: response.failed\ndata: {json.dumps({'type': 'response.failed', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'failed', 'model': payload.model, 'output': [], 'error': {'code': 502, 'message': _friendly_error(e)}}})}\n\n" + yield _sse( + "response.failed", + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": [], + "error": {"code": 502, "message": _friendly_error(e)}, + }, + }, + ) return if resp.status_code != 200: @@ -5304,7 +5832,24 @@ async def _responses_stream( resp.status_code, err_text[:500], ) - yield f"event: response.failed\ndata: {json.dumps({'type': 'response.failed', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'failed', 'model': payload.model, 'output': [], 'error': {'code': resp.status_code, 'message': f'llama-server error: {err_text[:500]}'}}})}\n\n" + yield _sse( + "response.failed", + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": [], + "error": { + "code": resp.status_code, + "message": f"llama-server error: {err_text[:500]}", + }, + }, + }, + ) return lines_iter = resp.aiter_lines() @@ -5334,17 +5879,38 @@ async def _responses_stream( continue delta = choices[0].get("delta", {}) or {} - content = delta.get("content") - if content: - full_text += content - delta_event = { - "type": "response.output_text.delta", - "item_id": msg_id, - "output_index": 0, - "content_index": 0, - "delta": content, - } - yield f"event: response.output_text.delta\ndata: {json.dumps(delta_event)}\n\n" + reasoning_delta, visible_delta = extractor.feed( + delta.get("content") or "", + delta.get("reasoning_content"), + ) + if reasoning_delta: + for event in _ensure_reasoning_open(): + yield event + full_reasoning += reasoning_delta + yield _sse( + "response.reasoning_text.delta", + { + "type": "response.reasoning_text.delta", + "item_id": reasoning_state["item_id"], + "output_index": reasoning_state["output_index"], + "content_index": 0, + "delta": reasoning_delta, + }, + ) + if visible_delta: + for event in _ensure_message_open(): + yield event + full_text += visible_delta + yield _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": visible_delta, + }, + ) for tc in delta.get("tool_calls") or []: idx = tc.get("index", 0) @@ -5354,14 +5920,13 @@ async def _responses_stream( # First chunk for this tool call -- allocate an # output_index and emit output_item.added. st = { - "output_index": next_output_index, + "output_index": _claim_output_index(), "item_id": f"fc_{uuid.uuid4().hex[:12]}", "call_id": tc.get("id") or "", "name": fn.get("name") or "", "arguments": "", "opened": False, } - next_output_index += 1 tool_call_state[idx] = st else: # Later chunks sometimes carry id/name only once; merge @@ -5384,7 +5949,7 @@ async def _responses_stream( "arguments": "", }, } - yield f"event: response.output_item.added\ndata: {json.dumps(item_added)}\n\n" + yield _sse("response.output_item.added", item_added) st["opened"] = True arg_delta = fn.get("arguments") or "" @@ -5396,7 +5961,7 @@ async def _responses_stream( "output_index": st["output_index"], "delta": arg_delta, } - yield f"event: response.function_call_arguments.delta\ndata: {json.dumps(args_delta_event)}\n\n" + yield _sse("response.function_call_arguments.delta", args_delta_event) elif arg_delta: # Buffer args until we can open the item (some models # send id/name in the same chunk as the first arg delta; @@ -5425,8 +5990,134 @@ async def _responses_stream( except Exception: pass - # ── Closing events for tool calls ── - for st in sorted(tool_call_state.values(), key = lambda s: s["output_index"]): + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + for event in _ensure_reasoning_open(): + yield event + full_reasoning += final_reasoning + yield _sse( + "response.reasoning_text.delta", + { + "type": "response.reasoning_text.delta", + "item_id": reasoning_state["item_id"], + "output_index": reasoning_state["output_index"], + "content_index": 0, + "delta": final_reasoning, + }, + ) + if final_visible: + for event in _ensure_message_open(): + yield event + full_text += final_visible + yield _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": final_visible, + }, + ) + if full_reasoning and not full_text and not tool_call_state: + for event in _ensure_message_open(): + yield event + full_text = full_reasoning + yield _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": full_text, + }, + ) + + close_items: list[tuple[int, str, dict[str, Any]]] = [] + if reasoning_state["opened"]: + close_items.append((reasoning_state["output_index"], "reasoning", reasoning_state)) + if message_state["opened"]: + close_items.append((message_state["output_index"], "message", message_state)) + close_items.extend((st["output_index"], "tool", st) for st in tool_call_state.values()) + + for _, kind, st in sorted(close_items, key = lambda item: item[0]): + if kind == "reasoning": + yield _sse( + "response.reasoning_text.done", + { + "type": "response.reasoning_text.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "text": full_reasoning, + }, + ) + yield _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "part": {"type": "reasoning_text", "text": full_reasoning}, + }, + ) + yield _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": st["output_index"], + "item": { + "type": "reasoning", + "id": st["item_id"], + "status": "completed", + "summary": [], + "content": [{"type": "reasoning_text", "text": full_reasoning}], + }, + }, + ) + continue + + if kind == "message": + yield _sse( + "response.output_text.done", + { + "type": "response.output_text.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "text": full_text, + }, + ) + yield _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "part": {"type": "output_text", "text": full_text, "annotations": []}, + }, + ) + yield _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": st["output_index"], + "item": { + "type": "message", + "id": st["item_id"], + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": full_text, "annotations": []} + ], + }, + }, + ) + continue + # If id/name never arrived (malformed upstream), synthesise so the # client still sees a coherent frame sequence. if not st["opened"]: @@ -5444,20 +6135,16 @@ async def _responses_stream( "arguments": "", }, } - yield f"event: response.output_item.added\ndata: {json.dumps(item_added)}\n\n" + yield _sse("response.output_item.added", item_added) if st["arguments"]: - yield ( - "event: response.function_call_arguments.delta\n" - "data: " - + json.dumps( - { - "type": "response.function_call_arguments.delta", - "item_id": st["item_id"], - "output_index": st["output_index"], - "delta": st["arguments"], - } - ) - + "\n\n" + yield _sse( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "item_id": st["item_id"], + "output_index": st["output_index"], + "delta": st["arguments"], + }, ) st["opened"] = True @@ -5468,7 +6155,7 @@ async def _responses_stream( "name": st["name"], "arguments": st["arguments"], } - yield f"event: response.function_call_arguments.done\ndata: {json.dumps(args_done)}\n\n" + yield _sse("response.function_call_arguments.done", args_done) item_done = { "type": "response.output_item.done", @@ -5482,14 +6169,7 @@ async def _responses_stream( "arguments": st["arguments"], }, } - yield f"event: response.output_item.done\ndata: {json.dumps(item_done)}\n\n" - - # ── Closing events for text message ── - yield f"event: response.output_text.done\ndata: {json.dumps({'type': 'response.output_text.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'text': full_text})}\n\n" - - yield f"event: response.content_part.done\ndata: {json.dumps({'type': 'response.content_part.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': {'type': 'output_text', 'text': full_text, 'annotations': []}})}\n\n" - - yield f"event: response.output_item.done\ndata: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': {'type': 'message', 'id': msg_id, 'status': 'completed', 'role': 'assistant', 'content': [{'type': 'output_text', 'text': full_text, 'annotations': []}]}})}\n\n" + yield _sse("response.output_item.done", item_done) # response.completed total_tokens = input_tokens + output_tokens @@ -5509,7 +6189,7 @@ async def _responses_stream( }, }, } - yield f"event: response.completed\ndata: {json.dumps(completed_response)}\n\n" + yield _sse("response.completed", completed_response) return StreamingResponse( event_generator(), @@ -5906,6 +6586,15 @@ async def anthropic_messages( ) if server_tools: + if bool(getattr(payload, "confirm_tool_calls", False)): + raise HTTPException( + status_code = 400, + detail = anthropic_error_body( + "confirm_tool_calls is not supported for Anthropic Messages server tools.", + status = 400, + err_type = "invalid_request_error", + ), + ) from core.inference.tools import ALL_TOOLS openai_tools = _select_anthropic_server_tools( diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index b30c0f5aea..3aae6f4209 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -31,6 +31,7 @@ class LlamaUpdateJob(BaseModel): from_tag: Optional[str] = None to_tag: Optional[str] = None error: Optional[str] = None + progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.") started_at: Optional[str] = None finished_at: Optional[str] = None @@ -40,7 +41,9 @@ class LlamaUpdateStatusResponse(BaseModel): False, description = "True when the install came from an Unsloth prebuilt (has a marker).", ) - update_available: bool = Field(False, description = "True when installed_tag != latest_tag.") + update_available: bool = Field( + False, description = "True when the latest release is genuinely newer than the install." + ) stale: bool = Field( False, description = "Update available AND install older than the staleness threshold." ) diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 3001c6b7c9..37d99a222e 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -10,12 +10,16 @@ from fastapi import APIRouter, Depends, HTTPException from auth.authentication import get_current_subject from core.inference.mcp_client import ( + TOOL_CACHE_INVALIDATING_FIELDS, + cache_tools, clear_oauth_tokens_async, + invalidate_tool_cache, is_stdio, list_tools_async, parse_server_headers, parse_stdio_command, probe_timeout, + record_probe_failure, stdio_mcp_enabled, ) from core.inference.mcp_config_import import parse_mcp_config @@ -198,6 +202,11 @@ async def update_mcp_server( ): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.update_server(server_id, changes) + # A new endpoint/auth makes cached tools wrong and disabling makes them + # unreachable, so drop them and let the next send re-probe; a rename + # leaves them valid. + if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS: + invalidate_tool_cache(server_id) return _row_to_response(mcp_servers_db.get_server(server_id)) @@ -209,6 +218,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c if old.get("use_oauth"): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.delete_server(server_id) + invalidate_tool_cache(server_id) @router.post("/{server_id}/refresh", response_model = McpServerProbeResult) @@ -238,8 +248,23 @@ async def refresh_mcp_server_tools( error = str(exc), exc_info = True, ) + current = mcp_servers_db.get_server(server_id) + if current is not None and not any( + current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + # Start the cool-off so the next chat send doesn't immediately re-hang + # on this server's timeout. If the row changed while the probe was + # awaiting, the failure belongs to the old config and must not park + # the newly edited server. + record_probe_failure(server_id, use_oauth) return McpServerProbeResult(ok = False, error = safe_curated_detail(exc)) + # Warm the chat-path cache so the next send skips re-probing. + current = mcp_servers_db.get_server(server_id) + if current is not None and not any( + current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + cache_tools(server_id, tools) return McpServerProbeResult(ok = True, tool_count = len(tools)) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 7151af33f3..a2f2eca81b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -51,6 +51,14 @@ def _is_hidden_model(*values: str | None) -> bool: return any(v and any(n in v.lower() for n in needles) for v in values) +def _safe_resolve(path: Path) -> Optional[str]: + """resolve() to a string, or None when the path is inaccessible.""" + try: + return str(path.resolve()) + except OSError: + return None + + backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) @@ -676,9 +684,9 @@ async def list_local_models( # trusted Path objects are used for FS access; the user string is # used for matching only, never for path construction. allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir] - if legacy_hf.is_dir(): + if _safe_is_dir(legacy_hf): allowed_roots.append(legacy_hf) - if hf_default.is_dir(): + if _safe_is_dir(hf_default): allowed_roots.append(hf_default) try: from utils.paths import studio_root, outputs_root @@ -702,15 +710,20 @@ async def list_local_models( try: local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + # Resolve once; an inaccessible aux cache must skip that scan, not 500. + hf_cache_real = _safe_resolve(hf_cache_dir) + legacy_real = _safe_resolve(legacy_hf) + default_real = _safe_resolve(hf_default) + # Scan legacy Unsloth HF cache for backward compatibility. - if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve(): + if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: local_models += _scan_hf_cache(legacy_hf) # Scan HF system default cache (may differ under env overrides). if ( - hf_default.is_dir() - and hf_default.resolve() != hf_cache_dir.resolve() - and hf_default.resolve() != legacy_hf.resolve() + _safe_is_dir(hf_default) + and default_real != hf_cache_real + and default_real != legacy_real ): local_models += _scan_hf_cache(hf_default) @@ -2069,13 +2082,10 @@ async def get_gguf_variants( best = _pick_best_gguf(filenames) default_variant = _extract_quant_label(best) if best else None - # Which variants are fully downloaded in the HF cache. For split - # GGUFs ALL shards must be present, so sum cached bytes per variant - # vs. the expected total. Cache dir casing may differ from the - # canonical repo_id, so match case-insensitively. - cached_bytes_by_quant: dict[str, int] = {} + # Per-snapshot so a split GGUF's shards must all sit in one snapshot; + # mmproj adapters are excluded so they can't inflate a quant's bytes. + cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = [] try: - import re as _re from huggingface_hub import constants as hf_constants if not _is_valid_repo_id(repo_id): @@ -2088,21 +2098,31 @@ async def get_gguf_variants( snapshots = entry / "snapshots" if snapshots.is_dir(): for snap in snapshots.iterdir(): + by_quant: dict[str, int] = {} for f in _iter_gguf_paths(snap): - q = _extract_quant_label(f.name) - cached_bytes_by_quant[q] = ( - cached_bytes_by_quant.get(q, 0) + f.stat().st_size - ) + if _is_mmproj_filename(f.name): + continue + try: + size = f.stat().st_size + except OSError: + continue # broken symlink / unreadable: skip + q = _extract_quant_label(f.name).lower() + by_quant[q] = by_quant.get(q, 0) + size + if by_quant: + cached_bytes_by_quant_per_snapshot.append(by_quant) break except Exception: pass def _is_fully_downloaded(variant) -> bool: - cached = cached_bytes_by_quant.get(variant.quant, 0) - if cached == 0 or variant.size_bytes == 0: + if variant.size_bytes == 0: return False - # Rounding tolerance (symlinks vs real sizes). - return cached >= variant.size_bytes * 0.99 + # Complete within one snapshot (tolerance for symlink size jitter). + quant = variant.quant.lower() + return any( + by_quant.get(quant, 0) >= variant.size_bytes * 0.99 + for by_quant in cached_bytes_by_quant_per_snapshot + ) return GgufVariantsResponse( repo_id = repo_id, @@ -2157,16 +2177,26 @@ async def get_gguf_download_progress( for entry in cache_dir.iterdir(): if entry.name.lower() == target: # Completed .gguf files for this variant in snapshots. + # Exclude mmproj so a vision adapter can't satisfy a same-label + # main variant (e.g. mmproj-F16 vs an F16 weight). for f in _iter_gguf_paths(entry): + if _is_mmproj_filename(f.name): + continue fname = f.name.lower().replace("-", "").replace("_", "") if not variant_lower or variant_lower in fname: - downloaded_bytes += f.stat().st_size + try: + downloaded_bytes += f.stat().st_size + except OSError: + continue # broken symlink / unreadable: skip # In-progress (.incomplete) downloads in blobs. blobs_dir = entry / "blobs" if blobs_dir.is_dir(): for f in blobs_dir.iterdir(): if f.is_file() and f.name.endswith(".incomplete"): - in_progress_bytes += f.stat().st_size + try: + in_progress_bytes += f.stat().st_size + except OSError: + continue break total_progress_bytes = downloaded_bytes + in_progress_bytes @@ -2300,11 +2330,22 @@ def _get_repo_size_cached(repo_id: str) -> int: def _all_hf_cache_scans(): - """scan_cache_dir results for the active, legacy, and default HF caches.""" + """scan_cache_dir for the active, legacy, and default HF caches. + + Each probe is isolated: an unreadable auxiliary cache (permission denied, + broken symlink, OS-redirected ~/.cache) is skipped, not fatal, so the + Downloaded list never blanks out and downloads never leak into Recommended. + """ from huggingface_hub import scan_cache_dir from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir - scans = [scan_cache_dir()] + scans = [] + # Guard the active cache too: degrade to "no downloads" instead of raising. + try: + scans.append(scan_cache_dir()) + except Exception as exc: + logger.warning("Could not scan active HF cache: %s", exc) + seen: set[str] = set() try: # Resolve the active cache dir for dedup. @@ -2314,13 +2355,18 @@ def _all_hf_cache_scans(): pass for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): - extra = extra_fn() - if extra.is_dir() and str(extra.resolve()) not in seen: - seen.add(str(extra.resolve())) - try: - scans.append(scan_cache_dir(cache_dir = str(extra))) - except Exception as exc: - logger.warning("Could not scan HF cache %s: %s", extra, exc) + try: + extra = extra_fn() + # is_dir()/resolve() can raise on an inaccessible path; skip it. + if not extra.is_dir(): + continue + resolved = str(extra.resolve()) + if resolved in seen: + continue + seen.add(resolved) + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc) return scans @@ -2379,6 +2425,38 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _blob_mtime(f) -> float: + """Blob modification time in epoch seconds (0.0 if unknown). + + Prefers HF metadata ``blob_last_modified``, falls back to stat(); uses + only mtimes (portable across Windows, macOS, Linux), never path parsing. + """ + ts = getattr(f, "blob_last_modified", None) + if isinstance(ts, (int, float)) and ts > 0: + return float(ts) + blob_path = getattr(f, "blob_path", None) + if blob_path: + try: + return float(Path(blob_path).stat().st_mtime) + except OSError: + pass + return 0.0 + + +def _repo_gguf_last_modified(repo_info) -> float: + """Newest mtime among a repo's primary (non-mmproj) GGUF blobs. + + Drives the Downloaded list's "last downloaded" ordering and groups a + multi-quant repo by its most recently downloaded quant. + """ + latest = 0.0 + for revision in repo_info.revisions: + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + latest = max(latest, _blob_mtime(f)) + return latest + + @router.get("/cached-gguf") async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): """List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" @@ -2399,17 +2477,30 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): continue key = repo_id.lower() existing = seen_lower.get(key) + last_modified = _repo_gguf_last_modified(repo_info) if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { + row = { "repo_id": repo_id, "size_bytes": total_size, "cache_path": str(repo_info.repo_path), } + # Keep the newest timestamp across duplicate caches; + # attach only when known so absent rows sort as oldest. + lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if lm > 0: + row["last_modified"] = lm + seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") continue - cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + # Newest download first; stable repo_id tie-break for equal/missing mtimes. + cached = sorted( + seen_lower.values(), + key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()), + ) return {"cached": cached} except Exception as e: logger.error(f"Error listing cached GGUF repos: {e}", exc_info = True) @@ -2447,18 +2538,39 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) ) if not has_weights: continue + last_modified = max( + ( + _blob_mtime(f) + for rev in repo_info.revisions + for f in rev.files + if f.file_name.endswith(_WEIGHT_EXTENSIONS) + ), + default = 0.0, + ) key = repo_id.lower() existing = seen_lower.get(key) if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { + row = { "repo_id": repo_id, "size_bytes": total_size, } + # Keep the newest timestamp across duplicate caches; + # attach only when known so absent rows sort as oldest. + lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if lm > 0: + row["last_modified"] = lm + seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") continue - cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + # Newest download first; stable repo_id tie-break for equal/missing mtimes. + cached = sorted( + seen_lower.values(), + key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()), + ) return {"cached": cached} except Exception as e: logger.error(f"Error listing cached models: {e}", exc_info = True) diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 64d1c3eab3..5a55c9b0bb 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -190,7 +190,8 @@ async def test_provider( """ Test connectivity to an external provider. - Makes a lightweight GET /models call to verify the API key works. + Makes a lightweight GET /models call to verify the API key works. Generic + custom endpoints use a chat-completions probe because /models is optional. encrypted_api_key is decrypted server-side and never stored. """ info = get_provider_info(payload.provider_type) @@ -212,6 +213,14 @@ async def test_provider( ) base_url = payload.base_url or info["base_url"] + if payload.provider_type == "custom": + if not base_url: + return ProviderTestResult( + success = False, + message = "Connection failed: Base URL is required for custom providers.", + models_count = None, + ) + client = ExternalProviderClient( provider_type = payload.provider_type, base_url = base_url, @@ -220,6 +229,26 @@ async def test_provider( ) try: + if payload.provider_type == "custom": + model_id = (payload.model_id or "").strip() + if not model_id: + return ProviderTestResult( + success = False, + message = "Connection failed: add a model ID to test custom providers.", + models_count = None, + ) + await client.chat_completion( + messages = [{"role": "user", "content": "ping"}], + model = model_id, + temperature = 0.0, + top_p = 1.0, + max_tokens = 1, + ) + return ProviderTestResult( + success = True, + message = "Connected successfully. Chat completions endpoint responded.", + models_count = None, + ) if info.get("model_list_mode") == "curated": await client.verify_models_endpoint_lightweight() return ProviderTestResult( diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d7687ffdee..281f03bcaf 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -125,6 +125,17 @@ async def start_training( backend = get_training_backend() + # S3 dataset loading needs the optional boto3 dependency. Reject early + # with a clear message so credentials are never accepted and then + # silently dropped on a host without boto3 installed. + if request.s3_config is not None: + from core.training.s3_dataset import boto3_available + if not boto3_available(): + raise HTTPException( + status_code = 501, + detail = "S3 dataset loading requires boto3. Install it with: pip install boto3", + ) + # Check before mutating state. if backend.is_training_active(): existing_job_id: Optional[str] = getattr(backend, "current_job_id", "") @@ -235,6 +246,7 @@ async def start_training( "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, "gpu_ids": request.gpu_ids, + "s3_config": request.s3_config.model_dump() if request.s3_config else None, } # Training page has no trust_remote_code toggle; as a safety net consult @@ -293,6 +305,10 @@ async def start_training( error = None, ) + except HTTPException: + # Deliberate rejections (S3 not implemented, resume validation) must + # reach the client with their original status, not a generic 500. + raise except ValueError as e: logger.warning("Rejected training GPU selection: %s", e) # Deliberate user-facing GPU-selection validation message. diff --git a/studio/backend/run.py b/studio/backend/run.py index 32dfe06a18..a883154c4a 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -232,6 +232,9 @@ def _verify_global_reachability(display_host: str, port: int) -> None: public internet. Synchronous so output lands between the banner URLs and the stop hint. Bounded at ~15s; failures swallowed (verifier failing != Studio failing). Only meaningful for a wildcard bind.""" + global _public_reachable + # Reset to "unknown" each run; set True/False only when the probe decides. + _public_reachable = None import ipaddress import json import time @@ -324,12 +327,14 @@ def _verify_global_reachability(display_host: str, port: int) -> None: print("", flush = True) if ok_nodes: + _public_reachable = True print( f"{ok_c} Reachability check: {url}/ is reachable from the " f"public internet ({ok_nodes}/{total} probe nodes connected).{reset}", flush = True, ) elif err_nodes: + _public_reachable = False print( f"{err_c} Reachability check: {url}/ is NOT reachable from " f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}", @@ -422,7 +427,9 @@ def _print_cloudflare_line() -> None: """Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up. Reads the module-level URL set by ``run_server``. Prints nothing when the - tunnel is disabled or failed -- failures are silently ignored. + tunnel is disabled or failed -- failures are silently ignored. When the public + reachability probe just failed (``_public_reachable is False``) but the tunnel + is up, reword to point the user at the Cloudflare link as the way in. """ if not _cloudflare_url: return @@ -430,7 +437,10 @@ def _print_cloudflare_line() -> None: accent = "\033[38;5;150;1m" reset = "\033[0m" - line = f" Secure link access via Cloudflare: {_cloudflare_url}" + if _public_reachable is False: + line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}" + else: + line = f" Secure link access via Cloudflare: {_cloudflare_url}" print(f"{accent}{line}{reset}" if stdout_supports_color() else line) @@ -622,6 +632,12 @@ _shutdown_event = None # None when there is no tunnel (loopback, disabled, or a silently-ignored failure). _cloudflare_url = None +# Public reachability from the last _verify_global_reachability run, read by the +# Cloudflare banner line. True when the public ip:port probe confirmed reachable, +# False when it confirmed NOT reachable, None when the probe did not run or could +# not decide (timeout, blocked, private address). +_public_reachable = None + _DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist" diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py new file mode 100644 index 0000000000..f66226b61d --- /dev/null +++ b/studio/backend/state/tool_approvals.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Per-call tool-call confirmation gate. + +When a chat request sets ``confirm_tool_calls``, the agentic loop pauses +before executing each tool and waits here for the user's decision, which +arrives via ``POST /api/inference/tool-confirm`` on a separate connection. + +Each gated call is identified by a unique ``approval_id`` (minted with +``new_approval_id``) that the loop both registers here and echoes in the +``tool_start`` stream event. The frontend sends that exact id back, so a +stale or duplicate confirmation -- or a second tool awaiting a decision in +the same session -- can never resolve the wrong call. ``session_id`` is +kept alongside purely as a scope check. + +The slot is registered with ``begin_tool_decision`` *before* the loop +yields ``tool_start``, closing the race where a fast confirmation (or an +auto "Always allow") could otherwise arrive before the waiter exists. +``wait_tool_decision`` then blocks and cleans up its own slot. +""" + +import secrets +import threading +from typing import Optional + +# Generous ceiling so a user can deliberate; cancellation (stop button / +# disconnect) still breaks the wait early via ``cancel_event``. +_DECISION_TIMEOUT = 3600.0 + +# Fed to the model as the tool result when the user denies a call, so it +# can adapt and keep responding instead of the turn ending abruptly. +TOOL_REJECTED_MESSAGE = "The user declined to run this tool call." + +_lock = threading.Lock() +# approval_id -> {"event": threading.Event, "decision": str|None, "session": str} +_pending: dict[str, dict] = {} + + +def new_approval_id() -> str: + """Mint an unguessable id for one pending tool-call confirmation.""" + return secrets.token_urlsafe(16) + + +def begin_tool_decision(session_id, approval_id) -> dict: + """Register a pending decision slot and return it. + + Call this *before* yielding the ``tool_start`` event so the waiter + always exists by the time the user's confirmation can arrive. + """ + slot = { + "event": threading.Event(), + "decision": None, + "session": session_id or "", + } + with _lock: + _pending[approval_id] = slot + return slot + + +def wait_tool_decision( + slot, + approval_id, + cancel_event = None, + timeout = _DECISION_TIMEOUT, +): + """Block on a slot from ``begin_tool_decision`` until the user decides. + + Returns ``"allow"`` or ``"deny"``. Falls back to ``"deny"`` if the wait + times out or generation is cancelled before the user decides. Always + removes its own slot on exit. + """ + try: + waited = 0.0 + while not slot["event"].wait(timeout = 0.5): + if cancel_event is not None and cancel_event.is_set(): + return "deny" + waited += 0.5 + if waited >= timeout: + return "deny" + return slot["decision"] or "deny" + finally: + with _lock: + if _pending.get(approval_id) is slot: + _pending.pop(approval_id, None) + + +def abort_tool_decision(slot, approval_id) -> None: + """Remove a slot that was announced but never entered ``wait_tool_decision``. + + Streaming wrappers may stop after ``tool_start`` is yielded and before + the loop resumes into ``wait_tool_decision``. In that case there is no + waiter to run the normal cleanup path, so the generator close path calls + this explicitly. + """ + with _lock: + if _pending.get(approval_id) is slot: + _pending.pop(approval_id, None) + + +def request_tool_decision( + session_id, + approval_id, + cancel_event = None, + timeout = _DECISION_TIMEOUT, +): + """Register and wait in one call (when the slot is not needed early).""" + slot = begin_tool_decision(session_id, approval_id) + return wait_tool_decision(slot, approval_id, cancel_event = cancel_event, timeout = timeout) + + +def resolve_tool_decision( + approval_id, + decision, + session_id = None, +) -> bool: + """Record the user's "allow"/"deny" decision and unblock the loop. + + Returns ``True`` if a pending call matched, ``False`` otherwise (e.g. a + stale or duplicate confirmation, or a session-scope mismatch). + + The first decision wins: once a slot's event is set, a later (duplicate or + out-of-order) confirmation for the same id is rejected without mutating the + recorded decision, so an Allow can never be flipped to Deny in the window + before the waiter reads ``slot["decision"]`` and pops the slot. + """ + if not approval_id: + return False + with _lock: + slot = _pending.get(approval_id) + if not slot: + return False + if session_id is not None and slot["session"] != (session_id or ""): + return False + if slot["event"].is_set(): + return False + slot["decision"] = decision + slot["event"].set() + return True diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index da8d9b5e66..85cfacbc27 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -652,7 +652,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at, r.ended_at, r.total_steps, r.final_step, r.final_loss, r.output_dir, r.duration_seconds, r.error_message, - r.loss_sparkline, r.display_name, + r.loss_sparkline, r.display_name, r.config_json, CASE WHEN r.status = 'stopped' AND r.output_dir IS NOT NULL diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index e1230ae113..f8d3f44d4b 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1543,6 +1543,19 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "tools" + def test_confirm_tool_calls_rejected_for_server_tools(self, monkeypatch): + backend = _mock_backend(monkeypatch) + payload = _basic_payload( + confirm_tool_calls = True, + tools = [{"type": "web_search_20250305", "name": "web_search"}], + ) + + with pytest.raises(HTTPException) as exc: + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert exc.value.status_code == 400 + assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"] + assert backend.calls == [] + def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch): backend = _mock_backend(monkeypatch) payload = _basic_payload( diff --git a/studio/backend/tests/test_apple_gpu_sensors.py b/studio/backend/tests/test_apple_gpu_sensors.py new file mode 100644 index 0000000000..50947d1380 --- /dev/null +++ b/studio/backend/tests/test_apple_gpu_sensors.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for Apple Silicon GPU sensors (SMC temperature + IOReport power).""" + +import ctypes +import platform +import time + +import pytest + +from utils.hardware import apple + +_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64" + + +class TestFourcc: + def test_roundtrip(self): + for key in ("#KEY", "Tg0D", "flt "): + assert apple._fourcc_str(apple._fourcc(key)) == key + + def test_known_value(self): + # "flt " FourCC, same constant macmon uses. + assert apple._fourcc("flt ") == 1718383648 + + +class TestWatts: + def test_millijoules(self): + assert apple._watts(2000, "mJ", 2.0) == pytest.approx(1.0) + + def test_microjoules(self): + assert apple._watts(5_000_000, "uJ", 1.0) == pytest.approx(5.0) + + def test_nanojoules(self): + assert apple._watts(1_500_000_000, "nJ", 1.0) == pytest.approx(1.5) + + def test_unknown_unit_returns_none(self): + assert apple._watts(1000, "J", 1.0) is None + + def test_zero_elapsed_returns_none(self): + assert apple._watts(1000, "mJ", 0.0) is None + + +class TestAverageValidTemps: + def test_averages_and_rounds(self): + assert apple._average_valid_temps([40.0, 50.0, 60.05]) == 50.0 + + def test_filters_invalid(self): + assert apple._average_valid_temps([-1.0, 0.0, 151.0, 42.0]) == 42.0 + + def test_empty_returns_none(self): + assert apple._average_valid_temps([]) is None + assert apple._average_valid_temps([0.0, 200.0]) is None + + +class TestSmcStructLayout: + def test_key_data_matches_smc_protocol_size(self): + # The AppleSMC user client rejects calls whose struct size differs. + assert ctypes.sizeof(apple._SMCKeyData) == 80 + + +@pytest.mark.skipif(not _IS_APPLE_SILICON, reason = "requires Apple Silicon") +class TestLiveSensors: + def test_gpu_temperature_in_plausible_range(self): + temp = apple.read_gpu_temperature_c() + assert temp is not None + assert 0.0 < temp <= 150.0 + + def test_gpu_power_after_baseline(self): + apple.read_gpu_power_w() # first call only sets the baseline + time.sleep(0.3) + power = apple.read_gpu_power_w() + assert power is not None + assert power >= 0.0 diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 20aa37a2b4..5a4ca68ab4 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -371,3 +371,156 @@ def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeyp "cache_path": str(vision_repo.repo_path), } ] + + +def _gfile(name: str, size: int, mtime: float) -> SimpleNamespace: + """A cached file carrying a Hugging Face ``blob_last_modified`` timestamp.""" + return SimpleNamespace( + file_name = name, + size_on_disk = size, + blob_path = None, + blob_last_modified = mtime, + ) + + +def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_path): + """An unreadable auxiliary cache (e.g. an inaccessible + ``~/.cache/huggingface/hub``) must be skipped, not abort the scan. + Regression guard for ``extra.is_dir()`` raising and wiping the response. + """ + import huggingface_hub + import utils.paths as paths_mod + + active = SimpleNamespace( + repos = [_repo("Org/Active", [_file("Q4_K_M.gguf", 5_000)], tmp_path / "active")] + ) + + def _fake_scan(cache_dir = None): + if cache_dir is None: + return active + raise AssertionError("auxiliary scan should have been skipped") + + class _Boom: + def is_dir(self): + raise PermissionError(13, "Permission denied") + + def resolve(self): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr(huggingface_hub, "scan_cache_dir", _fake_scan) + monkeypatch.setattr(paths_mod, "legacy_hf_cache_dir", lambda: _Boom()) + monkeypatch.setattr(paths_mod, "hf_default_cache_dir", lambda: _Boom()) + + scans = models_route._all_hf_cache_scans() + assert scans == [active] + + # End-to-end: the endpoint still returns the active cache's repo. + monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [active]) + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + assert result["cached"] == [ + { + "repo_id": "Org/Active", + "size_bytes": 5_000, + "cache_path": str(tmp_path / "active"), + } + ] + + +def test_list_cached_gguf_sorts_newest_first_grouping_by_latest_quant(monkeypatch, tmp_path): + """Downloaded is ordered newest-first, and a multi-quant repo is placed by + its most recently downloaded quant (``last_modified`` = newest quant).""" + older = _repo( + "Org/Older", + [_gfile("Older-Q4_K_M.gguf", 5_000, 1_000.0)], + tmp_path / "models--Org--Older", + ) + newer = _repo( + "Org/Newer", + [ + _gfile("Newer-Q4_K_M.gguf", 5_000, 2_000.0), + _gfile("Newer-Q8_0.gguf", 9_000, 3_000.0), # newest quant in the repo + ], + tmp_path / "models--Org--Newer", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [older, newer])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert [c["repo_id"] for c in result["cached"]] == ["Org/Newer", "Org/Older"] + assert result["cached"][0]["last_modified"] == 3_000.0 + assert result["cached"][1]["last_modified"] == 1_000.0 + + +def test_list_cached_gguf_dedupe_keeps_newest_timestamp(monkeypatch, tmp_path): + """Same repo in two caches with equal size keeps the newest last_modified, + regardless of scan order.""" + older = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 1_000.0)], tmp_path / "a") + newer = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 9_000.0)], tmp_path / "b") + for scans in ([older, newer], [newer, older]): # both orders + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda s = scans: [SimpleNamespace(repos = [s[0]]), SimpleNamespace(repos = [s[1]])], + ) + result = asyncio.run(models_route.list_cached_gguf(current_subject = "t")) + assert len(result["cached"]) == 1 + assert result["cached"][0]["last_modified"] == 9_000.0 + + +def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_path): + """The per-quant 'downloaded' flag is driven by the real weight file in a + single snapshot; an mmproj vision adapter (matching a quant label) must + not make that quant appear downloaded.""" + import huggingface_hub.constants as hf_constants + + variants = [ + SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000), + SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000), + ] + monkeypatch.setattr( + models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True) + ) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + + snap = tmp_path / "models--org--repo" / "snapshots" / "rev" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16" + + result = asyncio.run( + models_route.get_gguf_variants( + repo_id = "org/repo", hf_token = None, current_subject = "test-user" + ) + ) + + flags = {v.quant: v.downloaded for v in result.variants} + assert flags["Q4_K_M"] is True + assert flags["F16"] is False + + +def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path): + """A cached mmproj adapter must not count toward a same-label main + variant's download progress (mmproj-F16 vs an F16 weight).""" + import huggingface_hub.constants as hf_constants + + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + snap = tmp_path / "models--org--repo" / "snapshots" / "rev" + snap.mkdir(parents = True) + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # only the adapter on disk + + result = asyncio.run( + models_route.get_gguf_download_progress( + repo_id = "org/repo", + variant = "F16", + expected_bytes = 20_000, + current_subject = "test-user", + ) + ) + + assert result["downloaded_bytes"] == 0 + assert result["progress"] == 0 diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index 8208f7f83f..873547631d 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -13,6 +13,7 @@ import importlib.util import io import sys import tarfile +import types from pathlib import Path import pytest @@ -423,3 +424,56 @@ def test_run_server_gates_tunnel_on_wildcard(): source = _RUN_PY.read_text() assert "_cloudflare_enabled" in source assert 'host == "0.0.0.0"' in source + + +def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable): + """Exec the real _print_cloudflare_line source in isolation (run.py has heavy + deps), with the two module globals injected and startup_banner stubbed.""" + src = _RUN_PY.read_text() + tree = ast.parse(src) + func_src = next( + ast.get_source_segment(src, n) + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_print_cloudflare_line" + ) + stub = types.ModuleType("startup_banner") + stub.stdout_supports_color = lambda: False + monkeypatch.setitem(sys.modules, "startup_banner", stub) + captured: list[str] = [] + ns = { + "_cloudflare_url": cloudflare_url, + "_public_reachable": public_reachable, + "print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)), + } + exec(compile(func_src, "", "exec"), ns) + ns["_print_cloudflare_line"]() + return "\n".join(captured) + + +def test_cloudflare_line_reworded_when_public_unreachable(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = False + ) + assert "Use the secure link access via Cloudflare instead: https://x.trycloudflare.com" in out + + +def test_cloudflare_line_default_wording_when_reachable(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = True + ) + assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out + assert "Use the secure link" not in out + + +def test_cloudflare_line_default_wording_when_unknown(monkeypatch): + # Probe did not run / could not decide -> keep the existing wording. + out = _run_print_cloudflare_line( + monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = None + ) + assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out + assert "Use the secure link" not in out + + +def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch): + out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False) + assert out == "" diff --git a/studio/backend/tests/test_datacenter_gpu_tuning.py b/studio/backend/tests/test_datacenter_gpu_tuning.py new file mode 100644 index 0000000000..fd9b291e8a --- /dev/null +++ b/studio/backend/tests/test_datacenter_gpu_tuning.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data-center llama.cpp env tuning: FP32 accum (+ P2P / launch queues for +multi-GPU) must apply only to datacenter NVIDIA parts, never consumer GeForce, +AMD/ROCm, CPU or macOS. User values win; UNSLOTH_DISABLE_DC_TUNING=1 disables. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.llama_cpp import LlamaCppBackend + + +def _fake_torch( + names, + *, + hip = None, + cuda_ok = True, +): + """torch stub: version.hip, cuda.is_available/device_count, get_device_properties(i).name.""" + t = types.ModuleType("torch") + t.version = types.SimpleNamespace(hip = hip) + t.cuda = types.SimpleNamespace( + is_available = lambda: cuda_ok, + device_count = lambda: len(names), + get_device_properties = lambda i: types.SimpleNamespace(name = names[i]), + ) + return t + + +@pytest.fixture(autouse = True) +def _clear_cuda_visible_devices(monkeypatch): + """Detection reads CUDA_VISIBLE_DEVICES, so clear it by default (run unmasked, + physical id == ordinal) regardless of host; masked tests set it explicitly.""" + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + + +# --------------------------------------------------------------------------- +# _is_datacenter_gpu +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "names,expected", + [ + # Datacenter / professional parts. + (["NVIDIA A100-SXM4-80GB"], True), + (["NVIDIA A30"], True), + (["NVIDIA H100 80GB HBM3"], True), + (["NVIDIA H200"], True), + (["NVIDIA H800"], True), + (["NVIDIA GH200 480GB"], True), + (["NVIDIA B200"], True), + (["NVIDIA GB200"], True), + (["NVIDIA L40S"], True), + (["NVIDIA L4"], True), + (["NVIDIA RTX PRO 6000 Blackwell Server Edition"], True), + (["NVIDIA RTX 6000 Ada Generation"], True), + # Consumer GeForce: never. + (["NVIDIA GeForce RTX 4090"], False), + (["NVIDIA GeForce RTX 5090"], False), + (["NVIDIA GeForce RTX 3090"], False), + (["NVIDIA GeForce RTX 2080 Ti"], False), + (["NVIDIA GeForce GTX 1080"], False), + # Workstation/laptop: short markers must not match as substrings + # ("a100" in "A1000", "a30" in "A3000"). + (["NVIDIA RTX A1000 Laptop GPU"], False), + (["NVIDIA RTX A1000 6GB Laptop GPU"], False), + (["NVIDIA RTX A3000 Laptop GPU"], False), + # Homogeneous multi-DC: all must match. + (["NVIDIA B200", "NVIDIA B200"], True), + (["NVIDIA H100 80GB HBM3", "NVIDIA H100 80GB HBM3"], True), + # Mixed DC + consumer: non-DC, so tuning never lands on the GeForce. + (["NVIDIA B200", "NVIDIA GeForce RTX 4090"], False), + (["NVIDIA GeForce RTX 4090", "NVIDIA B200"], False), + ], +) +def test_is_datacenter_gpu(monkeypatch, names, expected): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(names)) + assert LlamaCppBackend._is_datacenter_gpu() is expected + + +def test_is_datacenter_gpu_respects_selection(monkeypatch): + # A mixed box where only the DC GPU is selected -> True; only consumer -> False. + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["NVIDIA B200", "NVIDIA GeForce RTX 4090"]), + ) + assert LlamaCppBackend._is_datacenter_gpu([0]) is True + assert LlamaCppBackend._is_datacenter_gpu([1]) is False + assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False + + +def test_is_datacenter_gpu_out_of_range_indices_skipped(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + # Out-of-range / negative indices are skipped; the one valid DC GPU still wins. + assert LlamaCppBackend._is_datacenter_gpu([0, 5, -1]) is True + # Only invalid indices -> nothing seen -> False (fail closed for the flag). + assert LlamaCppBackend._is_datacenter_gpu([5, 9]) is False + + +def test_is_datacenter_gpu_masked_host_physical_ids(monkeypatch): + # Mask 4,5,6,7 -> ordinals 0..3 == physical 4..7. PHYSICAL selection [4,5] + # must resolve, not index out of range (the pre-fix bug: 4 >= device_count). + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is True + assert LlamaCppBackend._is_datacenter_gpu([4, 5, 6, 7]) is True + assert LlamaCppBackend._is_datacenter_gpu(None) is True + assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False # not visible -> skip + + +def test_is_datacenter_gpu_masked_host_reordered(monkeypatch): + # Reordered mask preserves order: ordinal 0 -> physical 7, 1 -> 4, ... + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,4,5,6") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100 80GB HBM3"] * 4)) + assert LlamaCppBackend._is_datacenter_gpu([7, 4]) is True + + +def test_is_datacenter_gpu_masked_host_mixed_class(monkeypatch): + # Mask 4,5: physical 4 = GeForce, physical 5 = B200. Detection must follow the + # selected physical GPU, not a same-numbered ordinal. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["NVIDIA GeForce RTX 4090", "NVIDIA B200"]), + ) + assert LlamaCppBackend._is_datacenter_gpu([4]) is False + assert LlamaCppBackend._is_datacenter_gpu([5]) is True + assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is False + + +def test_is_datacenter_gpu_unparsable_mask_falls_back(monkeypatch): + # Unparsable (UUID) mask falls back to physical id == ordinal (mirrors + # _get_gpu_free_memory), so ordinal lookup still classifies the device. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-abcdef12") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + assert LlamaCppBackend._is_datacenter_gpu([0]) is True + + +def test_is_datacenter_gpu_rocm_is_false(monkeypatch): + # ROCm reuses torch.cuda.*; an MI300X must not qualify. + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["AMD Instinct MI300X"], hip = "6.2.0"), + ) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +def test_is_datacenter_gpu_no_cuda_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False)) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +def test_is_datacenter_gpu_missing_torch_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +# --------------------------------------------------------------------------- +# _effective_gpu_count +# --------------------------------------------------------------------------- + + +def test_effective_gpu_count_explicit_selection(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._effective_gpu_count([0]) == 1 + assert LlamaCppBackend._effective_gpu_count([0, 1, 2]) == 3 + + +def test_effective_gpu_count_none_uses_visible(monkeypatch): + # None -> visible device count. + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._effective_gpu_count(None) == 4 + + +def test_effective_gpu_count_no_cuda_is_zero(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False)) + assert LlamaCppBackend._effective_gpu_count(None) == 0 + + +def test_effective_gpu_count_missing_torch_is_zero(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._effective_gpu_count(None) == 0 + + +# --------------------------------------------------------------------------- +# _apply_datacenter_env (the env-injection decision) +# --------------------------------------------------------------------------- + + +def test_apply_env_single_dc_gpu_sets_only_fp32(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0]) is True + assert env == {"GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "1"} + assert "GGML_CUDA_P2P" not in env # no multi-GPU flags on one GPU + assert "CUDA_SCALE_LAUNCH_QUEUES" not in env + + +def test_apply_env_multi_dc_gpu_sets_all(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1" + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" + + +def test_apply_env_none_indices_uses_visible_count(monkeypatch): + # None on a 2x DC box -> multi-GPU flags applied. + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100", "NVIDIA H100"])) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, None) is True + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" + + +def test_apply_env_consumer_gpu_is_noop(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA GeForce RTX 4090"] * 2)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False + assert env == {} + + +def test_apply_env_user_value_wins(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2)) + env = { + "GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "0", # user explicitly disabled + "CUDA_SCALE_LAUNCH_QUEUES": "8x", # user override + } + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True + # setdefault must not clobber user values; the unset one still defaults. + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "0" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "8x" + assert env["GGML_CUDA_P2P"] == "1" + + +def test_apply_env_disable_flag_respected(monkeypatch): + monkeypatch.setenv("UNSLOTH_DISABLE_DC_TUNING", "1") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False + assert env == {} + + +def test_apply_env_fail_open_on_detection_error(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", None) # detection raises -> False + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0]) is False + assert env == {} + + +def test_apply_env_masked_host_multi_dc(monkeypatch): + # End-to-end masked host (mask 4,5,6,7, physical selection [4,5]): pre-fix + # applied no tuning; now all three multi-GPU flags must be set. + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [4, 5]) is True + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1" + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py new file mode 100644 index 0000000000..f333a590c2 --- /dev/null +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -0,0 +1,472 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import importlib.machinery +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +_BACKEND_DIR = Path(__file__).resolve().parent.parent + + +def _load_module( + module_name: str, + relative_path: str, + monkeypatch = None, +): + path = _BACKEND_DIR / relative_path + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + if monkeypatch is None: + sys.modules[module_name] = module + else: + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + +class _Router: + def get(self, *args, **kwargs): + return lambda fn: fn + + def post(self, *args, **kwargs): + return lambda fn: fn + + def delete(self, *args, **kwargs): + return lambda fn: fn + + +class _HTTPException(Exception): + def __init__( + self, + status_code: int, + detail: str | None = None, + ): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +class _LocalModelInfo: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +def _identity_decorator(*_args, **_kwargs): + return lambda fn: fn + + +def _install_lightweight_backend_stubs(monkeypatch): + fastapi = types.ModuleType("fastapi") + fastapi.APIRouter = lambda: _Router() + fastapi.Body = lambda default = None, **_kwargs: default + fastapi.Depends = lambda dependency = None, **_kwargs: dependency + fastapi.HTTPException = _HTTPException + fastapi.Query = lambda default = None, **_kwargs: default + fastapi.Request = object + monkeypatch.setitem(sys.modules, "fastapi", fastapi) + + fastapi_responses = types.ModuleType("fastapi.responses") + fastapi_responses.StreamingResponse = object + monkeypatch.setitem(sys.modules, "fastapi.responses", fastapi_responses) + + monkeypatch.setitem( + sys.modules, + "structlog", + types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ), + ) + loggers = types.ModuleType("loggers") + loggers.get_logger = lambda *args, **kwargs: _DummyLogger() + monkeypatch.setitem(sys.modules, "loggers", loggers) + + auth_pkg = types.ModuleType("auth") + auth_mod = types.ModuleType("auth.authentication") + auth_mod.get_current_subject = lambda: None + monkeypatch.setitem(sys.modules, "auth", auth_pkg) + monkeypatch.setitem(sys.modules, "auth.authentication", auth_mod) + + core_pkg = types.ModuleType("core") + core_export = types.ModuleType("core.export") + core_export.get_export_backend = lambda: None + core_inference = types.ModuleType("core.inference") + core_inference.get_inference_backend = lambda: None + monkeypatch.setitem(sys.modules, "core", core_pkg) + monkeypatch.setitem(sys.modules, "core.export", core_export) + monkeypatch.setitem(sys.modules, "core.inference", core_inference) + + utils_pkg = types.ModuleType("utils") + utils_pkg.__path__ = [] + utils_paths = types.ModuleType("utils.paths") + storage_roots = _load_module( + "utils.paths.storage_roots", + "utils/paths/storage_roots.py", + monkeypatch, + ) + utils_pkg.paths = utils_paths + utils_paths.storage_roots = storage_roots + utils_paths.is_local_path = lambda value: Path(str(value)).is_absolute() + utils_paths.outputs_root = lambda: Path("outputs") + utils_paths.exports_root = storage_roots.exports_root + utils_paths.resolve_cached_repo_id_case = lambda value: value + utils_paths.resolve_output_dir = lambda value = None: Path(value or "outputs") + utils_paths.resolve_export_dir = storage_roots.resolve_export_dir + monkeypatch.setitem(sys.modules, "utils", utils_pkg) + monkeypatch.setitem(sys.modules, "utils.paths", utils_paths) + + utils_utils = types.ModuleType("utils.utils") + utils_utils.log_and_http_error = lambda *args, **kwargs: (_ for _ in ()).throw( + _HTTPException(kwargs.get("status_code", 500), kwargs.get("detail")) + ) + utils_utils.safe_error_detail = lambda value: str(value) + monkeypatch.setitem(sys.modules, "utils.utils", utils_utils) + + utils_models = types.ModuleType("utils.models") + for name in ( + "scan_trained_models", + "scan_exported_models", + "scan_checkpoints", + "list_gguf_variants", + ): + setattr(utils_models, name, lambda *args, **kwargs: []) + for name in ( + "get_base_model_from_checkpoint", + "get_base_model_from_lora", + "load_model_defaults", + ): + setattr(utils_models, name, lambda *args, **kwargs: None) + utils_models.is_vision_model = lambda *args, **kwargs: False + utils_models.is_embedding_model = lambda *args, **kwargs: False + utils_models.ModelConfig = object + monkeypatch.setitem(sys.modules, "utils.models", utils_models) + + utils_model_config = types.ModuleType("utils.models.model_config") + utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None + utils_model_config._extract_quant_label = lambda value: value + utils_model_config.is_audio_input_type = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "utils.models.model_config", + utils_model_config, + ) + + models_pkg = types.ModuleType("models") + models_pkg.__path__ = [] + for name in ( + "CheckpointInfo", + "CheckpointListResponse", + "LocalModelListResponse", + "ModelCheckpoints", + "ModelDetails", + "LoRAScanResponse", + "LoRAInfo", + "ModelListResponse", + "LoadCheckpointRequest", + "ExportStatusResponse", + "ExportOperationResponse", + "ExportMergedModelRequest", + "ExportBaseModelRequest", + "ExportGGUFRequest", + "ExportLoRAAdapterRequest", + ): + setattr(models_pkg, name, object) + models_pkg.LocalModelInfo = _LocalModelInfo + monkeypatch.setitem(sys.modules, "models", models_pkg) + + models_models = types.ModuleType("models.models") + for name in ( + "BrowseEntry", + "BrowseFoldersResponse", + "GgufVariantDetail", + "GgufVariantsResponse", + "ScanFolderInfo", + "AddScanFolderRequest", + ): + setattr(models_models, name, object) + models_models.ModelType = str + monkeypatch.setitem(sys.modules, "models.models", models_models) + + models_responses = types.ModuleType("models.responses") + for name in ( + "LoRABaseModelResponse", + "VisionCheckResponse", + "EmbeddingCheckResponse", + ): + setattr(models_responses, name, object) + monkeypatch.setitem(sys.modules, "models.responses", models_responses) + + +def _install_pydantic_stub(monkeypatch): + pydantic = types.ModuleType("pydantic") + pydantic.BaseModel = object + pydantic.Field = lambda default = None, **_kwargs: default + pydantic.field_validator = _identity_decorator + monkeypatch.setitem(sys.modules, "pydantic", pydantic) + + +def _install_export_backend_stubs(monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + + unsloth = types.ModuleType("unsloth") + unsloth.FastLanguageModel = object + unsloth.FastVisionModel = object + unsloth._IS_MLX = True + unsloth.__spec__ = importlib.machinery.ModuleSpec("unsloth", loader = None) + monkeypatch.setitem(sys.modules, "unsloth", unsloth) + + unsloth_zoo = types.ModuleType("unsloth_zoo") + unsloth_zoo.__path__ = [] + unsloth_zoo.__spec__ = importlib.machinery.ModuleSpec( + "unsloth_zoo", + loader = None, + is_package = True, + ) + llama_cpp = types.ModuleType("unsloth_zoo.llama_cpp") + llama_cpp.LLAMA_CPP_DEFAULT_DIR = str(Path("/tmp/llama.cpp")) + llama_cpp._resolve_local_convert_script = lambda *args, **kwargs: None + llama_cpp.__spec__ = importlib.machinery.ModuleSpec( + "unsloth_zoo.llama_cpp", + loader = None, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo) + monkeypatch.setitem(sys.modules, "unsloth_zoo.llama_cpp", llama_cpp) + + huggingface_hub = types.ModuleType("huggingface_hub") + huggingface_hub.HfApi = object + huggingface_hub.ModelCard = object + monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub) + + utils_hardware = types.ModuleType("utils.hardware") + utils_hardware.clear_gpu_cache = lambda: None + monkeypatch.setitem(sys.modules, "utils.hardware", utils_hardware) + + utils_models = sys.modules["utils.models"] + utils_models.get_base_model_from_lora = lambda *args, **kwargs: None + utils_models.is_vision_model = lambda *args, **kwargs: False + + utils_model_config = sys.modules["utils.models.model_config"] + utils_model_config.detect_audio_type = lambda *args, **kwargs: None + + utils_paths = sys.modules["utils.paths"] + utils_paths.ensure_dir = lambda path: Path(path).mkdir(parents = True, exist_ok = True) + utils_paths.resolve_export_write_dir = lambda value = None: Path(value or "exports") + utils_paths.resolve_output_dir = lambda value = None: Path(value or "outputs") + + +def test_gguf_export_cleans_temp_dir_when_post_processing_fails(tmp_path, monkeypatch): + _install_export_backend_stubs(monkeypatch) + export_mod = _load_module("test_core_export_backend", "core/export/export.py", monkeypatch) + + cwd = tmp_path / "cwd" + save_dir = tmp_path / "export" + cwd.mkdir() + monkeypatch.chdir(cwd) + monkeypatch.setattr(export_mod, "resolve_export_write_dir", lambda _value: save_dir) + monkeypatch.setattr( + export_mod.shutil, + "move", + lambda *args, **kwargs: (_ for _ in ()).throw(OSError("move failed")), + ) + + class _Model: + def save_pretrained_gguf(self, model_save_path, tokenizer, quantization_method): + Path(model_save_path).mkdir(parents = True) + (Path(model_save_path) / "model.safetensors").write_bytes(b"weights") + (cwd / "converted.gguf").write_bytes(b"gguf") + + backend = export_mod.ExportBackend.__new__(export_mod.ExportBackend) + backend.current_model = _Model() + backend.current_tokenizer = object() + backend.current_checkpoint = None + + success, message, output_path = backend.export_gguf(str(save_dir), "Q4_K_M") + + assert success is False + assert "move failed" in message + assert output_path is None + assert list(save_dir.glob("_tmp_model_*")) == [] + + +def test_save_directory_validator_rejects_windows_parent_segments(monkeypatch): + _install_pydantic_stub(monkeypatch) + export_models = _load_module("test_models_export", "models/export.py", monkeypatch) + + with pytest.raises(ValueError, match = r"\.\."): + export_models._validate_save_directory(r"E:\AI\..\secret") + + +def test_save_directory_validator_allows_deep_absolute_paths(monkeypatch, tmp_path): + _install_pydantic_stub(monkeypatch) + export_models = _load_module("test_models_export_deep_path", "models/export.py", monkeypatch) + + deep_path = tmp_path + for index in range(40): + deep_path /= f"segment-{index:02d}" + raw = str(deep_path) + + assert len(raw) > 255 + assert export_models._validate_save_directory(raw) == raw + + +def test_save_directory_validator_rejects_long_path_component(monkeypatch, tmp_path): + _install_pydantic_stub(monkeypatch) + export_models = _load_module( + "test_models_export_long_component", "models/export.py", monkeypatch + ) + + with pytest.raises(ValueError, match = "path components"): + export_models._validate_save_directory(str(tmp_path / ("a" * 256))) + + +def test_export_write_dir_accepts_external_absolute_but_read_dir_rejects(tmp_path, monkeypatch): + storage_roots = _load_module( + "test_storage_roots_accept_external", + "utils/paths/storage_roots.py", + ) + + export_root = tmp_path / "exports" + external = tmp_path / "external" + export_root.mkdir() + external.mkdir() + monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root) + + assert storage_roots.resolve_export_write_dir(str(external)) == external + + with pytest.raises(ValueError, match = "path escapes root"): + storage_roots.resolve_export_dir(str(external)) + + +def test_export_write_dir_accepts_expanded_home_path(tmp_path, monkeypatch): + storage_roots = _load_module( + "test_storage_roots_accept_home_path", + "utils/paths/storage_roots.py", + ) + + export_root = tmp_path / "exports" + home = tmp_path / "home" + export_root.mkdir() + home.mkdir() + monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root) + if storage_roots.os.name == "nt": + monkeypatch.setenv("USERPROFILE", str(home)) + else: + monkeypatch.setenv("HOME", str(home)) + + assert storage_roots.resolve_export_write_dir("~/exports/model") == home / "exports" / "model" + + +def test_resolve_export_write_dir_rejects_backslash_parent_segment(): + storage_roots = _load_module( + "test_storage_roots_reject_parent", + "utils/paths/storage_roots.py", + ) + + with pytest.raises(ValueError, match = r"\.\."): + storage_roots.resolve_export_write_dir(r"exports\..\outside") + + +def test_export_write_dir_handles_non_native_windows_absolute_as_relative(tmp_path, monkeypatch): + storage_roots = _load_module( + "test_storage_roots_non_native_windows_path", + "utils/paths/storage_roots.py", + ) + + export_root = tmp_path / "exports" + export_root.mkdir() + monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root) + + if storage_roots.os.name == "nt": + pytest.skip("Windows drive paths are native on Windows") + + assert ( + storage_roots.resolve_export_write_dir(r"C:\exports\model") + == export_root / r"C:\exports\model" + ) + + +def test_export_details_registers_external_absolute_output(tmp_path, monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + export_route = _load_module( + "test_routes_export_external", + "routes/export.py", + monkeypatch, + ) + + output = tmp_path / "Gemma4_26B_gguf" + output.mkdir() + export_root = tmp_path / "studio" / "exports" + export_root.mkdir(parents = True) + registered = [] + + monkeypatch.setattr( + export_route, + "_try_register_external_export", + lambda path: (registered.append(path) is None, str(path)), + ) + monkeypatch.setattr( + "utils.paths.storage_roots.exports_root", + lambda: export_root, + ) + + details = export_route._export_details(str(output)) + + assert details == { + "output_path": str(output), + "scan_folder_registered": True, + "scan_folder_path": str(output), + } + assert registered == [output] + + +def test_export_details_does_not_register_contained_exports(tmp_path, monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + export_route = _load_module( + "test_routes_export_contained", + "routes/export.py", + monkeypatch, + ) + + export_root = tmp_path / "exports" + output = export_root / "model-gguf" + output.mkdir(parents = True) + + monkeypatch.setattr( + export_route, + "_try_register_external_export", + lambda path: pytest.fail(f"unexpected registration: {path}"), + ) + monkeypatch.setattr( + "utils.paths.storage_roots.exports_root", + lambda: export_root, + ) + + assert export_route._export_details(str(output)) == {"output_path": "model-gguf"} + + +def test_registered_absolute_export_folder_is_discoverable(tmp_path, monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + models_route = _load_module("test_routes_models", "routes/models.py", monkeypatch) + + export_dir = tmp_path / "Gemma4_26B_gguf" + export_dir.mkdir() + gguf_file = export_dir / "Gemma4_26B.BF16-00001-of-00002.gguf" + gguf_file.write_bytes(b"gguf") + + found = models_route._scan_models_dir(export_dir) + + assert len(found) == 1 + assert found[0].path == str(gguf_file) + assert found[0].source == "models_dir" diff --git a/studio/backend/tests/test_external_provider_usage_chunk.py b/studio/backend/tests/test_external_provider_usage_chunk.py index ebfd6a8f50..3ec9133718 100644 --- a/studio/backend/tests/test_external_provider_usage_chunk.py +++ b/studio/backend/tests/test_external_provider_usage_chunk.py @@ -144,6 +144,14 @@ def _make_openai_client() -> ExternalProviderClient: ) +def _make_custom_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "custom", + base_url = "http://custom.example/v1", + api_key = "", + ) + + def _anthropic_sse(events: list[dict]) -> bytes: chunks: list[str] = [] for event in events: @@ -180,6 +188,137 @@ def _usage_chunks(lines: list[str]) -> list[dict]: return out +def test_custom_provider_registry_is_hidden(): + from core.inference.providers import get_provider_info, list_available_providers + + info = get_provider_info("custom") + assert info is not None + assert info["hidden"] is True + assert "custom" not in {p["provider_type"] for p in list_available_providers()} + + +def test_custom_provider_uses_chat_completions_without_auth_key(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["headers"] = dict(request.headers) + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\ndata: [DONE]\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_custom_client() + lines = await _collect( + client.stream_chat_completion( + messages = [{"role": "user", "content": "ping"}], + model = "Qwen/Qwen3-0.6B", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + assert captured["url"] == "http://custom.example/v1/chat/completions" + assert "authorization" not in {k.lower() for k in captured["headers"]} + assert captured["body"]["model"] == "Qwen/Qwen3-0.6B" + assert any("ok" in line for line in lines) + + +def test_custom_provider_test_endpoint_probes_chat_completion(monkeypatch): + import importlib.util + import sys + from pathlib import Path + + module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" + spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + assert spec is not None + assert spec.loader is not None + providers_route = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = providers_route + spec.loader.exec_module(providers_route) + + captured: dict = {} + + class _FakeClient: + def __init__(self, **kwargs): + captured["init"] = kwargs + + async def chat_completion(self, **kwargs): + captured["chat_completion"] = kwargs + return {"choices": [{"message": {"content": "ok"}}]} + + async def list_models(self): + raise AssertionError("custom provider test must not call /models") + + async def close(self): + captured["closed"] = True + + monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient) + + async def run(): + return await providers_route.test_provider( + providers_route.ProviderTestRequest( + provider_type = "custom", + base_url = "http://custom.example/v1", + model_id = "Qwen/Qwen3-0.6B", + ), + current_subject = "unsloth", + ) + + result = _drive(run()) + assert result.success is True + assert result.models_count is None + assert captured["init"]["provider_type"] == "custom" + assert captured["chat_completion"]["model"] == "Qwen/Qwen3-0.6B" + assert captured["chat_completion"]["max_tokens"] == 1 + assert captured["closed"] is True + + +def test_custom_provider_test_endpoint_requires_model_id(monkeypatch): + import importlib.util + import sys + from pathlib import Path + + module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" + spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + assert spec is not None + assert spec.loader is not None + providers_route = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = providers_route + spec.loader.exec_module(providers_route) + + class _FakeClient: + def __init__(self, **kwargs): + pass + + async def close(self): + pass + + monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient) + + async def run(): + return await providers_route.test_provider( + providers_route.ProviderTestRequest( + provider_type = "custom", + base_url = "http://custom.example/v1", + ), + current_subject = "unsloth", + ) + + result = _drive(run()) + assert result.success is False + assert "model ID" in result.message + + def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch): sse_events = [ { diff --git a/studio/backend/tests/test_gemma4_chat_template_override.py b/studio/backend/tests/test_gemma4_chat_template_override.py new file mode 100644 index 0000000000..f726741aa5 --- /dev/null +++ b/studio/backend/tests/test_gemma4_chat_template_override.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-override of the chat template for ``unsloth/gemma-4-*-GGUF``. + +Studio ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking`` +defaulted off) and applies it to gemma-4 GGUF loads via the existing +``chat_template_override`` -> ``--chat-template-file`` path, so users do not need +to re-download quants. Pins the family matcher, the resolver precedence, the +bundled asset's reasoning/tool capabilities (which drive the "Preserve thinking" +UI toggle), the Jinja gate behaviour, and the reload-dedup interaction. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import pytest + +# ── chat_templates is dependency-light: load it directly so the pure-logic +# tests run without the studio venv / core.inference package side effects. ── +_CT_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_templates.py" +_ct_spec = importlib.util.spec_from_file_location("_gemma4_ct_test", _CT_PATH) +chat_templates = importlib.util.module_from_spec(_ct_spec) +_ct_spec.loader.exec_module(chat_templates) + +is_unsloth_gemma4_gguf = chat_templates.is_unsloth_gemma4_gguf +resolve_effective_chat_template_override = chat_templates.resolve_effective_chat_template_override +load_bundled_chat_template = chat_templates.load_bundled_chat_template +is_unsloth_gemma4_edge_gguf = chat_templates.is_unsloth_gemma4_edge_gguf + +BUNDLED = load_bundled_chat_template("gemma-4.jinja") # 12b / 26B-A4B / 31B +EDGE = load_bundled_chat_template("gemma-4-edge.jinja") # E2B / E4B + + +# ── Stubs so core.inference.llama_cpp imports without the full studio venv ── +def _stub_modules_ctx(): + """patch.dict context that stubs the heavy deps llama_cpp pulls in at import, + but only those NOT already importable (real httpx / structlog are kept when + present, e.g. in CI), and removes the stubs on exit so other tests are not + polluted.""" + from unittest.mock import patch + + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + _structlog_stub = _types.ModuleType("structlog") + _structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + overrides = { + name: stub + for name, stub in ( + ("loggers", _loggers_stub), + ("structlog", _structlog_stub), + ("httpx", _httpx_stub), + ) + if name not in sys.modules + } + return patch.dict(sys.modules, overrides) + + +def _detect_reasoning_flags(): + with _stub_modules_ctx(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags + + +# ── Family matcher ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "model_id,expected", + [ + ("unsloth/gemma-4-E2B-it-GGUF", True), + ("unsloth/gemma-4-E4B-it-GGUF", True), + ("unsloth/gemma-4-31B-it-GGUF", True), + ("unsloth/gemma-4-26B-A4B-it-GGUF", True), + ("UNSLOTH/GEMMA-4-E2B-IT-GGUF", True), # case-insensitive + ("gemma-4-E2B-it-GGUF", True), # owner-less shorthand -> unsloth/ + ("gemma-4-31B-it-GGUF", True), # owner-less shorthand -> unsloth/ + ("unsloth/gemma-4-E2B-it", False), # bf16, not GGUF + ("unsloth/gemma-3-4b-it-GGUF", False), # gemma 3 + ("google/gemma-4-31B-it-GGUF", False), # not unsloth + ("unsloth/Qwen3.5-9B-MTP-GGUF", False), + ("/home/user/models/gemma-4-E2B.Q4_K_M.gguf", False), # local path + ("", False), + (None, False), + ], +) +def test_is_unsloth_gemma4_gguf(model_id, expected): + assert is_unsloth_gemma4_gguf(model_id) is expected + + +# ── Resolver precedence ────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "model_id,expected_edge", + [ + ("unsloth/gemma-4-E2B-it-GGUF", True), + ("unsloth/gemma-4-E4B-it-GGUF", True), + ("UNSLOTH/GEMMA-4-E4B-IT-GGUF", True), + ("unsloth/gemma-4-12b-it-GGUF", False), + ("unsloth/gemma-4-26B-A4B-it-GGUF", False), + ("unsloth/gemma-4-31B-it-GGUF", False), + ("unsloth/gemma-3-4b-it-GGUF", False), + ], +) +def test_is_unsloth_gemma4_edge_gguf(model_id, expected_edge): + assert is_unsloth_gemma4_edge_gguf(model_id) is expected_edge + + +def test_resolver_returns_edge_template_for_e2b_e4b(): + for mid in ("unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF"): + out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None) + assert out == EDGE + assert out != BUNDLED + + +def test_resolver_handles_owner_less_shorthand(): + # ModelConfig.from_identifier prefixes unsloth/ for bare ids; the resolver + # runs before that, so it must apply the same normalization. + assert ( + resolve_effective_chat_template_override( + model_identifier = "gemma-4-E2B-it-GGUF", user_override = None + ) + == EDGE + ) + assert ( + resolve_effective_chat_template_override( + model_identifier = "gemma-4-31B-it-GGUF", user_override = None + ) + == BUNDLED + ) + + +def test_resolver_returns_standard_template_for_larger_models(): + for mid in ( + "unsloth/gemma-4-12b-it-GGUF", + "unsloth/gemma-4-26B-A4B-it-GGUF", + "unsloth/gemma-4-31B-it-GGUF", + ): + out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None) + assert out == BUNDLED + + +def test_resolver_user_override_wins(): + out = resolve_effective_chat_template_override( + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", user_override = "MY TEMPLATE" + ) + assert out == "MY TEMPLATE" + + +def test_resolver_blank_override_falls_back_to_bundled(): + out = resolve_effective_chat_template_override( + model_identifier = "unsloth/gemma-4-31B-it-GGUF", user_override = " " + ) + assert out == BUNDLED + + +def test_resolver_none_for_non_gemma(): + assert ( + resolve_effective_chat_template_override( + model_identifier = "unsloth/Llama-3.2-1B-Instruct-GGUF", user_override = None + ) + is None + ) + + +# ── Bundled asset content + capability classification ──────────────── + + +@pytest.mark.parametrize("tpl", [BUNDLED, EDGE]) +def test_bundled_template_has_preserve_thinking_defaulted_off(tpl): + assert "preserve_thinking" in tpl + assert "preserve_thinking | default(false)" in tpl + + +@pytest.mark.parametrize("name", ["gemma-4.jinja", "gemma-4-edge.jinja"]) +def test_bundled_templates_are_ascii(name): + # The temp file written for --chat-template-file must encode on any locale. + # Keeping the bundled templates ASCII avoids UnicodeEncodeError on non-UTF-8 + # Windows locales (cp932/cp1252) regardless of the writer's encoding. + text = load_bundled_chat_template(name) + non_ascii = sorted({c for c in text if ord(c) > 127}) + assert not non_ascii, f"{name} has non-ASCII chars: {non_ascii}" + + +@pytest.mark.parametrize("tpl", [BUNDLED, EDGE]) +def test_detect_reasoning_flags_on_bundled_template(tpl): + detect_reasoning_flags = _detect_reasoning_flags() + flags = detect_reasoning_flags(tpl, "unsloth/gemma-4-E2B-it-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking" + assert flags["reasoning_always_on"] is False + # This is what makes the "Preserve thinking" toggle appear in the UI. + assert flags["supports_preserve_thinking"] is True + assert flags["supports_tools"] is True + + +def test_edge_template_omits_empty_thought_block_on_thinking_off(): + """E2B/E4B must NOT emit the empty <|channel>thought block when + thinking is disabled; the larger-model template must. This is the only + intended difference between the two bundled templates.""" + EMPTY = "<|channel>thought\n" + msgs = [{"role": "user", "content": "hi"}] + edge_off = _render_with(EDGE, msgs, enable_thinking = False) + std_off = _render_with(BUNDLED, msgs, enable_thinking = False) + assert EMPTY not in edge_off, "edge (E2B/E4B) should not emit empty thought block" + assert EMPTY in std_off, "standard (12b/26B/31B) should emit empty thought block" + # With thinking ON neither appends the empty block at the prompt tail. + assert EMPTY not in _render_with(EDGE, msgs, enable_thinking = True) + + +# ── Jinja gate behaviour (off = omit prior reasoning, on = keep) ───── + + +def _render_with(tpl, messages, **kw): + pytest.importorskip("jinja2") # transitive via transformers; skip in minimal envs + from jinja2 import Environment, BaseLoader + + def raise_exception(msg): + raise RuntimeError(msg) + + env = Environment(loader = BaseLoader()) + return env.from_string(tpl).render( + messages = messages, + bos_token = "", + raise_exception = raise_exception, + add_generation_prompt = True, + **kw, + ) + + +def _render(messages, **kw): + return _render_with(BUNDLED, messages, **kw) + + +def _convo_with_prior_tool_reasoning(): + # Assistant tool-call turn with reasoning, BEFORE the last user message. + return [ + {"role": "user", "content": "q1"}, + { + "role": "assistant", + "reasoning_content": "SECRET_THOUGHT", + "tool_calls": [{"id": "c1", "function": {"name": "f", "arguments": {"x": 1}}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": "42"}, + {"role": "user", "content": "q2"}, + ] + + +def test_preserve_thinking_off_omits_prior_reasoning(): + # default(false): kwarg unset -> prior reasoning dropped before last user turn. + assert "SECRET_THOUGHT" not in _render(_convo_with_prior_tool_reasoning()) + + +def test_preserve_thinking_on_keeps_prior_reasoning(): + assert "SECRET_THOUGHT" in _render(_convo_with_prior_tool_reasoning(), preserve_thinking = True) + + +def test_enable_thinking_gates_think_token(): + assert "<|think|>" in _render([{"role": "user", "content": "hi"}], enable_thinking = True) + assert "<|think|>" not in _render([{"role": "user", "content": "hi"}], enable_thinking = False) + + +# ── Reload dedup interaction (why the route resolves the effective override) ── + + +def test_already_in_target_state_consistent_with_bundled_override(): + """The backend dedup compares the incoming override against the live one. + + The route resolves the bundled template up front so a re-load that omits + ``chat_template_override`` still matches (no spurious reload), while a raw + ``None`` would not. + """ + LlamaCppBackend = _import_backend() + + class _FakeProcess: + def terminate(self): ... + def wait(self, timeout = None): + return 0 + + def kill(self): ... + def poll(self): + return 0 + + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._healthy = True + backend._model_identifier = "unsloth/gemma-4-E2B-it-GGUF" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._speculative_type = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = BUNDLED # live server launched with the bundle + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + + common = dict( + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + extra_args = None, + is_vision = False, + ) + # Effective (resolved bundled) override -> already loaded, no reload. + assert backend._already_in_target_state(chat_template_override = BUNDLED, **common) is True + # Raw None (unresolved) -> false match, would force a needless reload. + assert backend._already_in_target_state(chat_template_override = None, **common) is False + + +def _import_backend(): + with _stub_modules_ctx(): + from core.inference.llama_cpp import LlamaCppBackend + return LlamaCppBackend diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py index 57bbc84fb0..7c6c514b8f 100644 --- a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py +++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py @@ -32,14 +32,19 @@ if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None: @pytest.fixture(autouse = True) def _clear_lemonade_release_cache(): - """Prevent cross-test pollution of the lemonade release lru_cache when - future tests vary the fetch_json mock return value.""" + """Prevent cross-test pollution of the lemonade release lru_cache and + selection-log dedup set when tests vary the fetch_json mock return value.""" _cache = getattr(_mod, "_fetch_lemonade_release_cached", None) + _logged: set | None = getattr(_mod, "_lemonade_selection_logged", None) if _cache is not None and hasattr(_cache, "cache_clear"): _cache.cache_clear() + if _logged is not None: + _logged.clear() yield if _cache is not None and hasattr(_cache, "cache_clear"): _cache.cache_clear() + if _logged is not None: + _logged.clear() _STUB_TAG = "b1262" diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index cb17e0d5e7..f90c4ba0e7 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -21,12 +21,25 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) + +class _NoopLogger: + """structlog-style logger: every method swallows positional + kwargs. + + A stdlib logging.Logger rejects structlog's keyword fields (e.g. + ``logger.warning(msg, error=...)``), which leaked into the update module's + error path and failed only when this file's stub loaded first. + """ + + def __getattr__(self, _name): + return lambda *a, **k: None + + _loggers_stub = _types.ModuleType("loggers") -_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +_loggers_stub.get_logger = lambda *a, **k: _NoopLogger() sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") -_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +_structlog_stub.get_logger = lambda *a, **k: _NoopLogger() sys.modules.setdefault("structlog", _structlog_stub) import pytest @@ -51,6 +64,11 @@ def _write_marker(install_dir: Path, **overrides) -> Path: .replace("+00:00", "Z"), } payload.update(overrides) + # The installer always writes `tag` and `release_tag` from the same release + # (a normalized base vs the full release tag), so keep the pair consistent + # when a test overrides only `tag`. + if "tag" in overrides and "release_tag" not in overrides: + payload["release_tag"] = overrides["tag"] install_dir.mkdir(parents = True, exist_ok = True) (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(payload)) return install_dir / "UNSLOTH_PREBUILT_INFO.json" @@ -303,3 +321,202 @@ def test_format_stale_warning_singular_day(): msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1}) assert "1 day" in msg assert "1 days" not in msg + + +# parse_base_build / is_behind. + + +def test_parse_base_build(): + assert fr.parse_base_build("b9596") == 9596 + assert fr.parse_base_build(" b9596 ") == 9596 + assert fr.parse_base_build("b9596-mix-e6f2453") == 9596 # mix suffix doesn't defeat it + assert fr.parse_base_build("9596") is None + assert fr.parse_base_build("master-abc") is None + assert fr.parse_base_build("") is None + assert fr.parse_base_build(None) is None + + +@pytest.mark.parametrize( + "installed, latest, expected", + [ + ( + "b9596-mix-e6f2453", + "b9596-mix-e6f2453", + False, + ), # already on the mix latest -> not behind + ("b9596", "b9594", False), # latest is an older build -> downgrade guard + ("b9596", "b9594-mix-xxx", False), # older mix latest -> still guarded + ("b9500", "b9596-mix-e6f2453", True), # newer base -> behind + ("b9596-mix-aaa", "b9596-mix-bbb", True), # new mix at same base -> behind + ("b9596", "b9596-mix-bbb", True), # clean -> mix at same base -> behind + ("b9596-mix-aaa", "b9596", False), # bare base never supersedes a mix install + ("b9596", "b9596", False), # identical -> not behind + (" b9596 ", "b9596", False), # whitespace-only diff -> not behind + ("master-abc", "master-def", True), # non-bNNNN both -> plain inequality + ("master-abc", "master-abc", False), + (None, "b9596", False), + ("b9596", None, False), + ], +) +def test_is_behind(installed, latest, expected): + assert fr.is_behind(installed, latest) is expected + + +def test_check_prebuilt_freshness_not_behind_on_mix_latest(monkeypatch, tmp_path): + # Installed the mix latest: marker base tag b9596, full release_tag with sha, + # GitHub latest is that same full tag. Must not report behind (sticky bug). + install_dir = tmp_path / "llama.cpp" + _write_marker(install_dir, tag = "b9596", release_tag = "b9596-mix-e6f2453") + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr( + fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + ) + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["behind"] is False + assert info["stale"] is False + + +def test_check_prebuilt_freshness_downgrade_guard(monkeypatch, tmp_path): + # A lagging latest (older build than installed) must never read as behind/stale. + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9585", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 30)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["behind"] is False + assert info["stale"] is False + + +def test_fetch_latest_release_tag_uses_publish_time(monkeypatch): + # Resolves newest by published_at (like the installer), skips drafts/prereleases, + # and does NOT just take GitHub's first/`/releases/latest` item. + import urllib.request + + class _Resp: + def __init__(self, payload): + self._p = json.dumps(payload).encode() + + def read(self): + return self._p + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + payload = [ + { + "tag_name": "b9518", + "draft": False, + "prerelease": False, + "published_at": "2026-06-04T21:11:19Z", + }, + { + "tag_name": "b9596-mix-e6f2453", + "draft": False, + "prerelease": False, + "published_at": "2026-06-11T22:50:41Z", + }, + { + "tag_name": "b9999-draft", + "draft": True, + "prerelease": False, + "published_at": "2026-06-12T00:00:00Z", + }, + ] + monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload)) + assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453" + + +# reset_caches(drop_disk=...) -- post-update stale same-base mix disk cache. + + +def _seed_disk_cache(tmp_path: Path, latest_tag: str) -> Path: + # Matches _cache_path_for under the fixture's stubbed _cache_dir. + cache_dir = tmp_path / ".freshness" + cache_dir.mkdir(exist_ok = True) + cache_file = cache_dir / "unslothai__llama.cpp.json" + cache_file.write_text(json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag})) + return cache_file + + +def test_reset_caches_drop_disk_removes_disk_cache(tmp_path): + cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa") + assert cache_file.exists() + fr.reset_caches(drop_disk = True) + assert not cache_file.exists() + + +def test_reset_caches_default_keeps_disk_cache(tmp_path): + # The no-arg form is in-memory only (its existing test-only contract); it + # must not delete the on-disk cache. + cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa") + fr.reset_caches() + assert cache_file.exists() + + +def test_reset_caches_drop_disk_on_missing_dir_is_noop(tmp_path): + # Fresh machine, no cache dir yet: drop_disk must be a quiet no-op. + assert not (tmp_path / ".freshness").exists() + fr.reset_caches(drop_disk = True) # must not raise + + +def test_drop_disk_lets_banner_fail_open_after_same_base_mix_swap(monkeypatch, tmp_path): + # P2 #2: the disk cache holds a still-fresh same-base mix (b9596-mix-aaa) + # from before an update to a *different* same-base mix (b9596-mix-bbb). + # The post-install path drops the disk cache; if the forced refresh is then + # offline, latest reads as None and the banner fails open -- instead of + # replaying the stale b9596-mix-aaa and falsely reading "behind". + _seed_disk_cache(tmp_path, "b9596-mix-aaa") + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9596", + release_tag = "b9596-mix-bbb", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + # GitHub unreachable for the rest of the test (the offline post-install + # refresh, and the later status check). + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + + fr.reset_caches(drop_disk = True) # exactly what the apply path now does + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["latest_tag"] is None + assert info["behind"] is False + assert info["stale"] is False + + +def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path): + # Contrast/guard for the case above: an in-memory-only reset leaves the + # stale same-base mix on disk, so an offline check replays it and falsely + # reads behind/stale. This is exactly the failure drop_disk removes; if a + # future change makes the no-arg reset also clear disk, the apply-path call + # and this guard should be revisited together. + _seed_disk_cache(tmp_path, "b9596-mix-aaa") + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9596", + release_tag = "b9596-mix-bbb", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + + fr.reset_caches() # in-memory only -> stale disk value survives + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["latest_tag"] == "b9596-mix-aaa" + assert info["behind"] is True + assert info["stale"] is True diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fa583ef53d..3c121c281d 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -21,6 +21,8 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) from core.inference.llama_cpp import LlamaCppBackend +from state import tool_approvals +from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision def _sse(delta: dict) -> str: @@ -70,6 +72,27 @@ def _tool_names(payload: dict) -> list[str]: ] +def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list[str]: + return [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": json.dumps(arguments), + }, + } + ] + } + ), + _done(), + ] + + def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): """llama-server may emit content first and then native delta.tool_calls. @@ -1149,3 +1172,151 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now.", "Final note after tool."] assert len(payloads) == 3 + + +def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch): + streams = [ + _structured_tool_call("python", {"code": "print(1)"}, "call_py"), + [_sse({"content": "Done."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: "approval-1") + monkeypatch.setattr( + "core.inference.llama_cpp.begin_tool_decision", + lambda *_a, **_k: object(), + ) + monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow") + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + ) + ) + + starts = [event for event in events if event.get("type") == "tool_start"] + assert len(starts) == 1 + assert starts[0]["approval_id"] + assert starts[0]["awaiting_confirmation"] is True + assert calls == [("python", {"code": "print(1)"})] + assert any(event.get("type") == "tool_end" and event.get("result") == "OK" for event in events) + + +def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch): + approval_id = "approval-close" + streams = [_structured_tool_call("python", {"code": "print(1)"}, "call_py")] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("tool should not run")), + ) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: approval_id) + + with tool_approvals._lock: + tool_approvals._pending.clear() + + gen = backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + ) + try: + assert next(gen)["type"] == "status" + start = next(gen) + assert start["type"] == "tool_start" + assert start["approval_id"] == approval_id + with tool_approvals._lock: + assert approval_id in tool_approvals._pending + finally: + gen.close() + + with tool_approvals._lock: + assert approval_id not in tool_approvals._pending + assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False + + +def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch): + streams = [[_sse({"content": "Done."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fail_autoinject(*_args, **_kwargs): + raise AssertionError("RAG autoinject must not run before approval") + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "use docs"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + rag_scope = {"thread_id": "t1"}, + ) + ) + + assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + + +def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypatch): + same_call = _structured_tool_call("python", {"code": "print(1)"}, "call_py") + streams = [ + same_call, + _structured_tool_call("python", {"code": "print(1)"}, "call_py_retry"), + [_sse({"content": "Done."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + decisions = iter(["deny", "allow"]) + approvals = iter(["approval-1", "approval-2"]) + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: next(approvals)) + monkeypatch.setattr( + "core.inference.llama_cpp.begin_tool_decision", + lambda *_a, **_k: object(), + ) + monkeypatch.setattr( + "core.inference.llama_cpp.wait_tool_decision", + lambda *_a, **_k: next(decisions), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 2, + confirm_tool_calls = True, + session_id = "sess", + ) + ) + + starts = [event for event in events if event.get("type") == "tool_start"] + ends = [event for event in events if event.get("type") == "tool_end"] + assert len(starts) == 2 + assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"] + assert calls == [("python", {"code": "print(1)"})] diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 05a107377c..3b8f511a7c 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -28,23 +28,75 @@ import utils.llama_cpp_update as upd # noqa: E402 MARKER = "UNSLOTH_PREBUILT_INFO.json" +class _FakeInstallerPopen: + """Stands in for the streamed installer process in _run_update.""" + + def __init__( + self, + cmd, + *, + returncode = 0, + lines = None, + on_start = None, + captured_kwargs = None, + **kwargs, + ): + if captured_kwargs is not None: + captured_kwargs.update(kwargs) + if on_start is not None: + on_start(list(cmd)) + self.returncode = returncode + self.stdout = iter(lines or []) + + def wait(self): + return self.returncode + + def kill(self): + pass + + +def _patch_installer_popen( + monkeypatch, + *, + returncode = 0, + lines = None, + on_start = None, + captured_kwargs = None, +): + monkeypatch.setattr( + upd.subprocess, + "Popen", + lambda cmd, **kw: _FakeInstallerPopen( + cmd, + returncode = returncode, + lines = lines, + on_start = on_start, + captured_kwargs = captured_kwargs, + **kw, + ), + ) + + def _write_install( dir_: Path, tag: str, repo: str = "unslothai/llama.cpp", asset: str | None = None, + release_tag: str | None = None, ) -> str: """Create a fake prebuilt install tree and return the llama-server path. ``asset`` is the bundle filename recorded in the marker; omit it to model an - older marker that predates asset-based ROCm forwarding (backward compat).""" + older marker that predates asset-based ROCm forwarding (backward compat). + ``release_tag`` is the full release tag (e.g. a ``b9596-mix-`` mix + build); defaults to ``tag`` for a plain prebuilt.""" bin_dir = dir_ / "build" / "bin" bin_dir.mkdir(parents = True, exist_ok = True) binary = bin_dir / "llama-server" binary.write_text("#!/bin/sh\necho stub\n") marker = { "tag": tag, - "release_tag": tag, + "release_tag": release_tag or tag, "published_repo": repo, "installed_at_utc": "2020-01-01T00:00:00Z", "bundle_profile": "cuda13-newer", @@ -57,10 +109,13 @@ def _write_install( @pytest.fixture(autouse = True) -def _clean_state(monkeypatch): +def _clean_state(monkeypatch, tmp_path): freshness.reset_caches() upd._reset_job_for_tests() upd._resolve_memo.clear() + # Isolate the freshness disk cache so the suite never writes the real + # ~/.unsloth cache (the default when storage_roots can't be imported). + monkeypatch.setattr(freshness, "_cache_dir", lambda: tmp_path / ".freshness_cache") # Deterministic markerless paths: no host-pinned binary, no custom dir. monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False) monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) @@ -260,14 +315,15 @@ def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path): def _fake_run(cmd, **kwargs): cmd = list(cmd) - # Status polls probe `llama-server --version`; keep the installer argv. - if "--version" in cmd: - return _Proc() - captured["cmd"] = cmd - _write_install(install_dir, "b9585") # installer writes the marker + assert "--version" in cmd # only status polls still use run() return _Proc() + def _on_start(cmd): + captured["cmd"] = cmd + _write_install(install_dir, "b9585") # installer writes the marker + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = _on_start) res = upd.start_update() assert res["started"] is True, res @@ -298,21 +354,27 @@ def test_start_update_happy_path(monkeypatch, tmp_path): stdout = "installed" stderr = "" - def _fake_run(cmd, **kwargs): - cmd = list(cmd) - # Status polls probe `llama-server --version`; keep the installer argv. - if "--version" in cmd: - return _Proc() + def _on_start(cmd): captured["cmd"] = cmd # Simulate the installer writing a new marker with the latest tag. _write_install(install_dir, "b9518") - return _Proc() - monkeypatch.setattr(upd.subprocess, "run", _fake_run) + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = [ + "[llama-prebuilt] resolving release\n", + "Downloading llama.zip: 35.0% (12.0 MiB/35.0 MiB) at 9.0 MiB/s\n", + "Downloading llama.zip: 80.0% (28.0 MiB/35.0 MiB) at 9.0 MiB/s\n", + ], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) res = upd.start_update() assert res["started"] is True assert res["job"]["from_tag"] == "b9493" + assert res["job"]["progress"] == 0.0 # Wait for the background worker. deadline = time.time() + 10 @@ -328,6 +390,10 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert str(install_dir) in captured["cmd"] assert "--llama-tag" in captured["cmd"] and "latest" in captured["cmd"] assert "unslothai/llama.cpp" in captured["cmd"] + # Progress lines were parsed and success pins progress at 1.0. + assert job["progress"] == 1.0 + # The worker asks the installer for fine-grained progress milestones. + assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): @@ -335,13 +401,9 @@ def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): binary = _write_install(install_dir, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") - class _Proc: - returncode = 2 - stdout = "" - stderr = "boom: network error" - - monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + _patch_installer_popen(monkeypatch, returncode = 2, lines = ["boom: network error\n"]) res = upd.start_update() assert res["started"] is True @@ -410,14 +472,15 @@ def _capture_install_cmd( def _fake_run(cmd, **kwargs): cmd = list(cmd) - # Status polls probe `llama-server --version`; keep the installer argv. - if "--version" in cmd: - return _Proc() - captured["cmd"] = cmd - _write_install(install_dir, latest, repo = repo, asset = asset) + assert "--version" in cmd # only status polls still use run() return _Proc() + def _on_start(cmd): + captured["cmd"] = cmd + _write_install(install_dir, latest, repo = repo, asset = asset) + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = _on_start) res = upd.start_update() assert res["started"] is True, res @@ -535,18 +598,12 @@ def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path): seen = {} - class _Proc: - returncode = 0 - stdout = "ok" - stderr = "" - - def _fake_run(cmd, **kwargs): + def _on_start(cmd): # The maintenance flag must be set while the installer runs. seen["flag_during_install"] = backend._llama_update_in_progress _write_install(install_dir, "b9518") - return _Proc() - monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = _on_start) res = upd.start_update() assert res["started"] is True @@ -567,16 +624,12 @@ def test_update_clears_maintenance_flag_on_installer_failure(monkeypatch, tmp_pa binary = _write_install(install_dir, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") backend = _FakeBackend() _inject_backend(monkeypatch, backend) - class _Proc: - returncode = 1 - stdout = "" - stderr = "boom" - - monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + _patch_installer_popen(monkeypatch, returncode = 1, lines = ["boom\n"]) res = upd.start_update() assert res["started"] is True @@ -606,16 +659,7 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): monkeypatch.setitem(sys.modules, "routes", routes_pkg) monkeypatch.setitem(sys.modules, "routes.inference", inference_mod) - class _Proc: - returncode = 0 - stdout = "ok" - stderr = "" - - def _fake_run(cmd, **kwargs): - _write_install(install_dir, "b9518") - return _Proc() - - monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = lambda cmd: _write_install(install_dir, "b9518")) res = upd.start_update() assert res["started"] is True @@ -759,3 +803,41 @@ def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path): res = upd.start_update() assert res["started"] is False assert res["reason"] == "up_to_date" + + +# --- mix-tag detection + apply guard (the reported banner bug) --- + + +def test_status_not_offered_on_mix_latest(monkeypatch, tmp_path): + # Installed the mix latest; GitHub latest is that same full tag -> no banner. + binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + ) + st = upd.get_update_status() + assert st["update_available"] is False + assert st["installed_tag"] == "b9596" + assert st["latest_tag"] == "b9596-mix-e6f2453" + + +def test_status_not_offered_when_latest_lags(monkeypatch, tmp_path): + # A lagging latest (older build than installed) must never be offered. + binary = _write_install(tmp_path / "llama.cpp", "b9585") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + st = upd.get_update_status() + assert st["update_available"] is False + + +def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path): + # A direct POST / stale banner must not reinstall when already on the latest. + binary = _write_install(tmp_path / "llama.cpp", "b9596", release_tag = "b9596-mix-e6f2453") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9596-mix-e6f2453" + ) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "up_to_date" diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 09775707ee..f3a3ea1ec4 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -25,8 +25,12 @@ _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag parse_cache_override = _lsa.parse_cache_override parse_ctx_override = _lsa.parse_ctx_override +parse_split_mode_override = _lsa.parse_split_mode_override resolve_cache_type_kv = _lsa.resolve_cache_type_kv +resolve_tensor_parallel = _lsa.resolve_tensor_parallel strip_shadowing_flags = _lsa.strip_shadowing_flags +strip_split_mode_only = _lsa.strip_split_mode_only +extra_args_disable_mmproj = _lsa.extra_args_disable_mmproj validate_extra_args = _lsa.validate_extra_args @@ -509,6 +513,128 @@ def test_strip_shadowing_flags_defaults_strip_everything(): assert out == [] +# ── --split-mode (Tensor Parallelism toggle) ───────────────────────── +# Soft-shadowed exactly like --cache-type-*: pass-through allowed (keeps +# the row/none/layer modes the boolean toggle doesn't expose), stripped +# on inherit, and reconciled back into the round-tripped tensor_parallel +# state. + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode", "tensor"], + ["--split-mode", "row"], + ["--split-mode", "none"], + ["--split-mode", "layer"], + ["-sm", "tensor"], + ["--split-mode=row"], + ["-sm=tensor"], + ], +) +def test_split_mode_passes_through(args): + # Not denylisted -- a user keeps row/none/layer via extras. + assert validate_extra_args(args) == args + + +def test_split_mode_is_not_managed(): + assert is_managed_flag("--split-mode") is False + assert is_managed_flag("-sm") is False + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--split-mode", "tensor"], "tensor"), + (["--split-mode", "row"], "row"), + (["-sm", "none"], "none"), + (["--split-mode=layer"], "layer"), + (["-sm=tensor"], "tensor"), + # last-wins when supplied twice + (["-sm", "row", "--split-mode", "tensor"], "tensor"), + ], +) +def test_parse_split_mode_override(args, expected): + assert parse_split_mode_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode"], + ["-sm"], + ["--split-mode", "-c", "4096"], # next token is a flag, not a value + ], +) +def test_parse_split_mode_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "split-mode|'-sm'"): + parse_split_mode_override(args) + + +def test_validate_extra_args_rejects_malformed_split_mode(): + # Validation catches a value-less --split-mode at the boundary, + # mirroring the early --ctx-size / --cache-type checks. + with pytest.raises(ValueError, match = "split-mode"): + validate_extra_args(["--split-mode"]) + + +@pytest.mark.parametrize( + "args,fallback,expected", + [ + # No override -> fall back to the toggle value, both directions. + (["--top-k", "20"], True, True), + (["--top-k", "20"], False, False), + (None, True, True), + ([], False, False), + # Explicit override wins: tensor -> on, anything else -> off, + # regardless of the toggle fallback. + (["--split-mode", "tensor"], False, True), + (["-sm", "tensor"], False, True), + (["--split-mode", "row"], True, False), + (["--split-mode", "none"], True, False), + (["--split-mode", "layer"], True, False), + (["--split-mode=tensor"], False, True), + # Case-insensitive on the mode string. + (["--split-mode", "TENSOR"], False, True), + # last-wins across multiple --split-mode flags. + (["-sm", "tensor", "--split-mode", "row"], True, False), + ], +) +def test_resolve_tensor_parallel(args, fallback, expected): + assert resolve_tensor_parallel(args, fallback) is expected + + +def test_strip_shadowing_flags_drops_split_mode_when_requested(): + out = strip_shadowing_flags( + ["--split-mode", "row", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) + assert out == ["--top-k", "20"] + + +def test_extra_args_disable_mmproj_detects_flag(): + assert extra_args_disable_mmproj(["--no-mmproj"]) is True + assert extra_args_disable_mmproj(["--threads", "12", "--no-mmproj"]) is True + assert extra_args_disable_mmproj(["--no-mmproj-auto"]) is True + + +def test_extra_args_disable_mmproj_false_when_absent(): + assert extra_args_disable_mmproj(None) is False + assert extra_args_disable_mmproj(["--threads", "12"]) is False + + +def test_extra_args_disable_mmproj_last_wins(): + assert extra_args_disable_mmproj(["--no-mmproj", "--mmproj-auto"]) is False + assert extra_args_disable_mmproj(["--mmproj-auto", "--no-mmproj-auto"]) is True + + def test_strip_shadowing_flags_drops_model_draft_with_spec(): # --model-draft (and aliases) are Studio-managed since the separate # MTP drafter support: an inherited copy must not last-wins-override @@ -523,6 +649,86 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec(): assert out == ["--top-k", "20"] +def test_strip_shadowing_flags_keeps_split_mode_when_not_requested(): + # No tensor_parallel field supplied on the Apply -> an inherited + # --split-mode survives (mirrors the chat-template keep behavior). + out = strip_shadowing_flags( + ["--split-mode", "row", "--top-k", "20"], + strip_context = True, + strip_cache = True, + strip_spec = True, + strip_template = True, + strip_split_mode = False, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + +def test_strip_shadowing_flags_drops_split_mode_short_alias_and_equals(): + assert strip_shadowing_flags(["-sm", "tensor", "--top-k", "20"], strip_split_mode = True) == [ + "--top-k", + "20", + ] + assert strip_shadowing_flags(["--split-mode=row", "--seed", "-1"], strip_split_mode = True) == [ + "--seed", + "-1", + ] + + +def test_strip_shadowing_flags_defaults_strip_split_mode_too(): + # The route's already-loaded comparator (no kwargs) must see a stored + # --split-mode as a shadowing flag so it forces a reload. + assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode", "tensor", "-c", "4096"], + ["-sm", "tensor", "-c", "4096"], + ["--split-mode=tensor", "-c", "4096"], + ["-sm=tensor", "-c", "4096"], + ], +) +def test_strip_split_mode_only_keeps_other_shadow_flags(args): + # Every --split-mode form (long/short, space/=) is dropped; -c survives. + assert strip_split_mode_only(args) == ["-c", "4096"] + + +def test_strip_split_mode_only_preserves_none_and_empty(): + # None means "inherit"; [] means "explicit empty" -- both must round-trip. + assert strip_split_mode_only(None) is None + assert strip_split_mode_only([]) == [] + + +def test_strip_shadowing_flags_drops_tensor_split_with_split_mode(): + # --tensor-split is coupled to the split mode: stripped together so a stale + # ratio can't override Studio's computed tensor split. Other flags survive. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_tensor_split_when_not_requested(): + # strip_split_mode=False keeps the whole split group (mode + ratios). + assert strip_shadowing_flags( + ["--tensor-split", "1,1", "--top-k", "20"], strip_split_mode = False + ) == ["--tensor-split", "1,1", "--top-k", "20"] + + +def test_strip_split_mode_only_drops_tensor_split_too(): + # Downgrade / layer fallback must drop the coupled --tensor-split (all forms). + assert strip_split_mode_only( + ["--split-mode", "tensor", "--tensor-split", "1,1", "-c", "4096"] + ) == ["-c", "4096"] + assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index ede3cf15d4..90b1ade03c 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -598,3 +598,614 @@ def test_safetensors_agentic_empty_allowlist_still_means_allow_all(): ) # Empty allow-list = run anything (preserved contract). assert calls == [("python", {"code": "1"})] or len(calls) >= 1 + + +# ── discovery cache ───────────────────────────────────────────────── + + +def _one_tool(name = "echo"): + return [{"name": name, "inputSchema": {"type": "object", "properties": {}}}] + + +def test_get_enabled_mcp_tools_caches_discovery(tmp_path, monkeypatch): + """A second send must serve tools from cache instead of re-probing.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + calls: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + calls.append(url) + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + first = asyncio.run(tools_mod.get_enabled_mcp_tools()) + second = asyncio.run(tools_mod.get_enabled_mcp_tools()) + + assert len(calls) == 1 # probed once, cache hit on the second send + assert [t["function"]["name"] for t in first] == ["mcp__s1__echo"] + assert first == second + + +def test_get_enabled_mcp_tools_does_not_cache_failures(tmp_path, monkeypatch): + """A failed probe isn't cached: once the cool-off elapses, it's retried.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + attempts = {"n": 0} + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + attempts["n"] += 1 + if attempts["n"] == 1: + raise RuntimeError("server down") + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # failure -> empty + # Expire the cool-off (an until-time in the past) so the server is retried. + mcp_client._probe_cooloff_until["s1"] = 0.0 + second = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert attempts["n"] == 2 # retried after the cool-off, not cached + assert [t["function"]["name"] for t in second] == ["mcp__s1__echo"] + + +def test_refresh_warms_tool_cache(tmp_path, monkeypatch): + """Clicking Refresh must populate the cache the chat path reads.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake_refresh( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + return _one_tool() + + monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok and res.tool_count == 1 + + def boom(*a, **k): + raise AssertionError("chat path re-probed despite a warm cache") + + monkeypatch.setattr(tools_mod, "list_tools_async", boom) + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert [t["function"]["name"] for t in specs] == ["mcp__s1__echo"] + + +def test_update_url_evicts_tool_cache(tmp_path, monkeypatch): + """Re-pointing the URL must drop the old endpoint's cached tools.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("stale")}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "https://new/mcp"), current_subject = "u" + ) + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch): + """A rename touches no endpoint, so the cache must survive it.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + cached = _one_tool() + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": cached}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u") + ) + assert mcp_client.get_cached_tools("s1") == cached + + +def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch): + """Disabling a server must drop its cached tools, not leave them unread.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(is_enabled = False), current_subject = "u") + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_delete_evicts_tool_cache(tmp_path, monkeypatch): + """Deleting a server must not leave its tools cached.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + asyncio.run(routes_mcp.delete_mcp_server("s1", current_subject = "u")) + assert mcp_client.get_cached_tools("s1") is None + + +def test_invalidate_tool_cache_clears_all(monkeypatch): + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {"a": _one_tool(), "b": _one_tool()}) + mcp_client.invalidate_tool_cache() + assert mcp_client.get_cached_tools("a") is None + assert mcp_client.get_cached_tools("b") is None + + +def test_get_enabled_mcp_tools_probes_only_uncached(tmp_path, monkeypatch): + """An already-cached server must not be re-probed alongside a cold one.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("cached")}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", is_enabled = True) + mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://b/mcp", is_enabled = True) + + probed: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + probed.append(url) + return _one_tool("fresh") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert probed == ["https://b/mcp"] # only the uncached server is probed + assert sorted(t["function"]["name"] for t in specs) == ["mcp__s1__cached", "mcp__s2__fresh"] + + +def test_get_enabled_mcp_tools_partial_failure_caches_healthy(tmp_path, monkeypatch): + """One server failing must not stop the others from being cached/served.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://bad/mcp", is_enabled = True) + mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://good/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + if "bad" in url: + raise RuntimeError("down") + return _one_tool("ok") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert [t["function"]["name"] for t in specs] == ["mcp__s2__ok"] + assert mcp_client.get_cached_tools("s1") is None # failure not cached + assert mcp_client.get_cached_tools("s2") == _one_tool("ok") # healthy cached + + +def test_get_enabled_mcp_tools_caches_empty_tool_list(tmp_path, monkeypatch): + """A server exposing zero tools is cached as [] (a hit), not re-probed.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + calls: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + calls.append(url) + return [] + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert len(calls) == 1 # [] is a cache hit, not re-probed every send + assert mcp_client.get_cached_tools("s1") == [] + + +def test_update_headers_evicts_tool_cache(tmp_path, monkeypatch): + """Changing auth headers must drop tools discovered under the old headers.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate(headers = {"Authorization": "Bearer new"}), + current_subject = "u", + ) + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A config edit landing during an in-flight probe must not be clobbered + by the now-stale probe result (TOCTOU on the cache write).""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # Simulate a PUT landing while we are awaiting the probe. + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert specs == [] # stale result is neither served... + assert mcp_client.get_cached_tools("s1") is None # ...nor cached + + +def test_get_enabled_mcp_tools_no_cooloff_when_config_changes_mid_failed_probe( + tmp_path, monkeypatch +): + """An edit landing while a probe of the OLD config is failing must not park + a cool-off on the now-fresh config -- else the re-pointed server the user + just fixed is needlessly skipped for the whole cool-off window.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # The user re-points the server while the old endpoint's probe fails. + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + raise RuntimeError("old endpoint down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + # The failure was for the OLD config, so the new one must stay re-probable. + assert not mcp_client.in_failure_cooloff("s1") + + +def test_get_enabled_mcp_tools_no_cooloff_when_server_deleted_mid_failed_probe( + tmp_path, monkeypatch +): + """A delete landing while a probe fails must not leave an orphan cool-off + entry keyed by the since-removed server id.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.delete_server("s1") + raise RuntimeError("down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert "s1" not in mcp_client._probe_cooloff_until # no orphan cool-off + + +def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff(tmp_path, monkeypatch): + """A down server is probed once, then skipped during the cool-off instead + of being re-probed (and re-hung) on every send.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + attempts = {"n": 0} + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + attempts["n"] += 1 + raise RuntimeError("down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # probes, fails + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # within cool-off + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # still skipped + assert attempts["n"] == 1 # only the first send probed + + +def test_cache_tools_clears_failure_cooloff(monkeypatch): + """A successful probe lifts a server's failure cool-off.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_client.record_probe_failure("s1") + assert mcp_client.in_failure_cooloff("s1") + mcp_client.cache_tools("s1", _one_tool()) + assert not mcp_client.in_failure_cooloff("s1") + + +def test_oauth_failure_cools_off_longer_than_plain(monkeypatch): + """An OAuth server's failure cools off longer than a plain server's, so its + multi-minute probe hang doesn't recur every minute.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_client.record_probe_failure("plain", use_oauth = False) + mcp_client.record_probe_failure("oauth", use_oauth = True) + assert mcp_client._probe_cooloff_until["oauth"] > mcp_client._probe_cooloff_until["plain"] + + +def test_invalidate_clears_failure_cooloff(monkeypatch): + """Eviction drops the failure cool-off so an edited server re-probes at once.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {"s1": 1.0, "s2": 2.0}) + mcp_client.invalidate_tool_cache("s1") + assert "s1" not in mcp_client._probe_cooloff_until + assert "s2" in mcp_client._probe_cooloff_until + mcp_client.invalidate_tool_cache() + assert mcp_client._probe_cooloff_until == {} + + +def test_refresh_failure_records_cooloff(tmp_path, monkeypatch): + """A failed manual refresh starts the cool-off so the next chat send does + not immediately hang on the down server.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + raise RuntimeError("down") + + monkeypatch.setattr(routes_mcp, "list_tools_async", boom) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok is False + assert mcp_client.in_failure_cooloff("s1") + + +def test_refresh_drops_result_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A manual refresh must not warm the chat cache with tools discovered + under an old config if the server is edited while the probe is in flight.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake_refresh( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + return _one_tool("stale") + + monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok and res.tool_count == 1 + assert mcp_client.get_cached_tools("s1") is None + + +def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A manual refresh failure for an old config must not cool off the freshly + edited server.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + raise RuntimeError("old endpoint down") + + monkeypatch.setattr(routes_mcp, "list_tools_async", boom) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok is False + assert not mcp_client.in_failure_cooloff("s1") + + +def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe(tmp_path, monkeypatch): + """A delete landing while a probe is in flight must drop the now-orphan + result -- the `fresh is None` arm of the mid-probe TOCTOU guard. The + result is neither served nor cached under the since-removed id.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # Simulate a DELETE landing while we await the probe. + mcp_servers_db.delete_server("s1") + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert specs == [] # orphan result not served + assert mcp_client.get_cached_tools("s1") is None # nor cached under a gone id + + +def test_oauth_probe_failure_in_chat_path_uses_long_cooloff(tmp_path, monkeypatch): + """When an OAuth server fails discovery during a send, the chat path must + record the OAuth (long) cool-off, not the plain one -- otherwise its + multi-minute browser hang recurs every minute.""" + import asyncio + import time + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "https://x/mcp", + is_enabled = True, + use_oauth = True, + ) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + raise RuntimeError("oauth down") + + monkeypatch.setattr(tools_mod, "list_tools_async", boom) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert mcp_client.in_failure_cooloff("s1") + # The recorded window must exceed the plain cool-off, proving the OAuth + # branch (use_oauth=True) fired -- not the 60 s default. + remaining = mcp_client._probe_cooloff_until["s1"] - time.monotonic() + assert remaining > mcp_client.FAILED_PROBE_COOLOFF_SECONDS diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py index c6a4898d5d..9a3e8d6882 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -19,6 +19,10 @@ from storage import mcp_servers_db def _reset_db(tmp_path, monkeypatch): monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + # The discovered-tool cache is process-global and keyed by server id; tests + # reuse "stdio1", so clear it (and the failure cool-off) for isolation — + # otherwise a prior test's warm cache makes discovery skip its probe. + mcp_client.invalidate_tool_cache() def _enable(monkeypatch): diff --git a/studio/backend/tests/test_native_context_length.py b/studio/backend/tests/test_native_context_length.py index 0290bb1308..de1ca0649e 100644 --- a/studio/backend/tests/test_native_context_length.py +++ b/studio/backend/tests/test_native_context_length.py @@ -271,6 +271,7 @@ class TestPydanticModels: def test_load_response_has_field(self): """Field exists in LoadResponse.model_fields.""" assert "native_context_length" in LoadResponse.model_fields + assert "context_length" in LoadResponse.model_fields def test_load_response_defaults_none(self): """Omitting native_context_length defaults to None.""" @@ -319,6 +320,7 @@ class TestPydanticModels: def test_status_response_has_field(self): """Field exists in InferenceStatusResponse.model_fields.""" assert "native_context_length" in InferenceStatusResponse.model_fields + assert "context_length" in InferenceStatusResponse.model_fields def test_status_response_has_chat_template_field(self): """Status includes chat_template so the UI can rehydrate after refresh.""" @@ -347,6 +349,18 @@ class TestPydanticModels: roundtripped = LoadResponse.model_validate_json(resp.model_dump_json()) assert roundtripped.native_context_length == 131072 + def test_context_length_roundtrip(self): + """Runtime context_length serializes for non-GGUF/hub models.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + context_length = 8192, + ) + roundtripped = LoadResponse.model_validate_json(resp.model_dump_json()) + assert roundtripped.context_length == 8192 + # ===================================================================== # D. TestRouteCompleteness -- source-level verification @@ -408,6 +422,16 @@ class TestRouteCompleteness: "native_context_length" not in block ), f"Non-GGUF LoadResponse should not set native_context_length:\n{block[:200]}" + def test_non_gguf_load_responses_set_runtime_context_length(self): + """Non-GGUF LoadResponse blocks report runtime context_length.""" + blocks = self._find_construction_blocks("LoadResponse") + non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b] + assert non_gguf, "Expected at least one non-GGUF LoadResponse block" + for block in non_gguf: + assert ( + "context_length" in block + ), f"Non-GGUF LoadResponse should set context_length:\n{block[:200]}" + def test_status_path(self): """InferenceStatusResponse construction with llama_backend has the field.""" blocks = self._find_construction_blocks("InferenceStatusResponse") @@ -420,6 +444,21 @@ class TestRouteCompleteness: found ), "No InferenceStatusResponse block with llama_backend has native_context_length" + def test_non_gguf_status_path_reports_runtime_context_length(self): + """Non-GGUF InferenceStatusResponse reports context_length from model_info.""" + blocks = self._find_construction_blocks("InferenceStatusResponse") + found = False + for block in blocks: + if "is_gguf = False" in block and "context_length" in block: + found = True + break + assert found, "No non-GGUF InferenceStatusResponse block with context_length" + + def test_openai_models_listing_reports_context_length(self): + """/v1/models includes context_length when the backend knows it.""" + assert 'entry["context_length"]' in self._source + assert 'model_info.get("context_length")' in self._source + # ===================================================================== # E. TestEdgeCases diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 1d994b46c0..baa3e50b95 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -28,11 +28,13 @@ from models.inference import ( ChatMessage, CompletionChoice, CompletionMessage, + ResponsesRequest, ) from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from routes.inference import ( + _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, _clamp_finish_reason, @@ -401,6 +403,28 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) + def test_confirm_tool_calls_rejected_for_provider_tools(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "external_model": "gpt-4.1", + "enable_tools": True, + "enabled_tools": ["web_search"], + "confirm_tool_calls": True, + }, + ) + + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["param"] == "confirm_tool_calls" + assert "only supported for local streaming tools" in body["error"]["message"] + def test_logprobs_rejected_until_supported(self, monkeypatch): class _UnusedBackend: is_loaded = False @@ -480,6 +504,7 @@ class TestChatCompletionRequestToolFields: def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False + supports_tools = False class _InferenceBackend: active_model_name = "test-model" @@ -495,6 +520,45 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) + def test_confirm_tool_calls_requires_streaming_for_safetensors_tools(self, monkeypatch): + import routes.inference as inference_route + + class _NoGGUFBackend: + is_loaded = False + supports_tools = False + + class _InferenceBackend: + active_model_name = "test-model" + models = {"test-model": {"chat_template_info": {"template": "chatml"}}} + + def generate_chat_completion_with_tools(self, **kwargs): + raise AssertionError("tool loop should be rejected before starting") + + def generate_chat_completion(self, **kwargs): + raise AssertionError("plain path should not be used") + + monkeypatch.setattr( + inference_route, + "_detect_safetensors_features", + lambda backend, chat_template: {"supports_tools": True}, + ) + client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "enable_tools": True, + "enabled_tools": ["web_search"], + "confirm_tool_calls": True, + "stream": False, + }, + ) + + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["param"] == "confirm_tool_calls" + assert "requires stream=true" in body["error"]["message"] + def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( messages = [ @@ -724,6 +788,22 @@ class TestPassthroughReasoningKwargs: ) assert body["chat_template_kwargs"] == {"reasoning_effort": "high"} + def test_reasoning_effort_none_forwarded_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False, reasoning_effort = "none"), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "none"} + + def test_reasoning_effort_minimal_maps_to_low_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = True, reasoning_effort = "minimal"), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "low"} + def test_enable_thinking_maps_to_effort_for_effort_style_models(self): body = _build_openai_passthrough_body( self._payload(enable_thinking = False), @@ -1206,6 +1286,45 @@ class TestGgufVisionToolRouting: assert captured["kwargs"]["disable_parallel_tool_use"] is True + def test_confirm_tool_calls_requires_streaming_for_gguf_tools(self, monkeypatch): + import routes.inference as inf_mod + + def _plain(**kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**kwargs): + raise AssertionError("tool loop should be rejected before starting") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + model_identifier = "test-gguf", + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + enable_tools = True, + enabled_tools = ["web_search"], + confirm_tool_calls = True, + stream = False, + messages = [{"role": "user", "content": "search once"}], + ) + + with pytest.raises(HTTPException) as exc: + self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert exc.value.status_code == 400 + assert "requires stream=true" in exc.value.detail["error"]["message"] + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1295,3 +1414,47 @@ class TestGgufVisionToolRouting: assert seen_seeds == expected assert [choice["index"] for choice in body["choices"]] == [0, 1, 2] + + +# ===================================================================== +# Responses API -> Chat Completions translation: chat_template_kwargs +# (e.g. {"enable_thinking": true}) sent via the Responses extra-body must +# reach the built ChatCompletionRequest's typed ``enable_thinking`` field, +# otherwise /v1/responses silently ignores reasoning control (issue #6198). +# ===================================================================== + + +class TestResponsesChatTemplateKwargs: + _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + + def test_enable_thinking_lifted_from_extra_body(self): + payload = ResponsesRequest( + model = "qwen-local", + input = "What is 100 - 67?", + chat_template_kwargs = {"enable_thinking": True}, + ) + chat_req = _build_chat_request(payload, self._messages, stream = False) + assert chat_req.enable_thinking is True + + def test_enable_thinking_false_lifted_from_extra_body(self): + payload = ResponsesRequest( + model = "qwen-local", + input = "hi", + chat_template_kwargs = {"enable_thinking": False}, + ) + chat_req = _build_chat_request(payload, self._messages, stream = True) + assert chat_req.enable_thinking is False + + def test_no_chat_template_kwargs_leaves_enable_thinking_unset(self): + payload = ResponsesRequest(model = "qwen-local", input = "hi") + chat_req = _build_chat_request(payload, self._messages, stream = False) + assert chat_req.enable_thinking is None + + def test_chat_template_kwargs_without_enable_thinking_is_ignored(self): + payload = ResponsesRequest( + model = "qwen-local", + input = "hi", + chat_template_kwargs = {"some_other_flag": True}, + ) + chat_req = _build_chat_request(payload, self._messages, stream = False) + assert chat_req.enable_thinking is None diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 69fb0a78c2..ae7ff729bd 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -36,6 +36,7 @@ import json import httpx import pytest +from fastapi.responses import JSONResponse from pydantic import ValidationError from models.inference import ( @@ -46,6 +47,7 @@ from models.inference import ( ResponsesInputMessage, ResponsesOutputFunctionCall, ResponsesOutputMessage, + ResponsesOutputReasoning, ResponsesOutputTextContent, ResponsesOutputTextPart, ResponsesRequest, @@ -59,6 +61,7 @@ from routes.inference import ( _chat_tool_calls_to_responses_output, _normalise_responses_input, _responses_tool_output_text, + _responses_non_streaming, _responses_stream, _translate_responses_tool_choice_to_chat, _translate_responses_tools_to_chat, @@ -284,6 +287,59 @@ class TestBuildChatRequest: assert chat_req.parallel_tool_calls is False + def test_chat_template_kwargs_enable_thinking_true_is_lifted(self): + payload = ResponsesRequest( + input = "hi", + chat_template_kwargs = {"enable_thinking": True}, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.enable_thinking is True + + def test_chat_template_kwargs_enable_thinking_false_is_lifted(self): + payload = ResponsesRequest( + input = "hi", + chat_template_kwargs = {"enable_thinking": False}, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.enable_thinking is False + + def test_reasoning_effort_high_enables_local_thinking(self): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.reasoning_effort == "high" + assert chat_req.enable_thinking is True + + def test_reasoning_effort_none_disables_local_thinking(self): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"}) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.reasoning_effort == "none" + assert chat_req.enable_thinking is False + + def test_explicit_enable_thinking_false_disables_reasoning_effort(self): + payload = ResponsesRequest( + input = "hi", + reasoning = {"effort": "high"}, + chat_template_kwargs = {"enable_thinking": False}, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.reasoning_effort == "none" + assert chat_req.enable_thinking is False + # ===================================================================== # _normalise_responses_input — multi-turn tool mapping @@ -544,6 +600,119 @@ class TestChatToolCallsToResponsesOutput: assert items[0]["arguments"] == "" +# ===================================================================== +# Non-streaming Responses adapter +# ===================================================================== + + +class TestResponsesNonStreamingAdapter: + class _Request: + pass + + @staticmethod + def _run_with_message( + monkeypatch, + message, + payload = None, + llama_backend = None, + ): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + return JSONResponse( + content = { + "model": "test-model", + "choices": [{"message": message}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + ) + + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + if llama_backend is not None: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: llama_backend) + payload = payload or ResponsesRequest(input = "hi") + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_non_streaming( + payload, messages, TestResponsesNonStreamingAdapter._Request() + ) + return json.loads(response.body.decode()) + + return asyncio.run(run()) + + def test_think_block_becomes_reasoning_item_before_message(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + body = self._run_with_message( + monkeypatch, + {"content": "plan33"}, + payload = payload, + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][0]["summary"] == [] + assert body["output"][1]["content"][0]["text"] == "33" + assert "" not in body["output"][1]["content"][0]["text"] + assert "" not in body["output"][1]["content"][0]["text"] + + def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch): + body = self._run_with_message(monkeypatch, {"content": "show x tags"}) + + assert [item["type"] for item in body["output"]] == ["message"] + assert body["output"][0]["content"][0]["text"] == "show x tags" + + def test_non_reasoning_gguf_keeps_literal_think_tags_visible(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + body = self._run_with_message( + monkeypatch, + {"content": "show x tags"}, + payload = payload, + llama_backend = SimpleNamespace( + is_loaded = True, + reasoning_always_on = False, + supports_reasoning = False, + ), + ) + + assert [item["type"] for item in body["output"]] == ["message"] + assert body["output"][0]["content"][0]["text"] == "show x tags" + + def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch): + body = self._run_with_message( + monkeypatch, + { + "content": "33", + "reasoning_content": [ + {"type": "reasoning_text", "text": "plan"}, + {"type": "reasoning_text", "text": " next"}, + ], + }, + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan next"}] + assert body["output"][1]["content"][0]["text"] == "33" + + def test_plain_content_remains_message_only(self, monkeypatch): + body = self._run_with_message(monkeypatch, {"content": "33"}) + + assert [item["type"] for item in body["output"]] == ["message"] + assert body["output"][0]["content"][0]["text"] == "33" + + def test_reasoning_only_is_also_visible_message_text(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + body = self._run_with_message( + monkeypatch, + {"content": "plan"}, + payload = payload, + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"][0]["text"] == "plan" + assert body["output"][1]["content"][0]["text"] == "plan" + + # ===================================================================== # Streaming Responses adapter # ===================================================================== @@ -570,6 +739,262 @@ class TestResponsesStreamAdapter: if line.startswith(prefix) ] + @staticmethod + def _install_stream_mock( + monkeypatch, + chunks, + *, + supports_reasoning = True, + reasoning_always_on = False, + ): + import routes.inference as inf_mod + + def handler(request: httpx.Request) -> httpx.Response: + content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def _client(*args, **kwargs): + return real_async_client( + transport = transport, + timeout = kwargs.get("timeout", 600), + ) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + context_length = 4096, + base_url = "http://llama.test", + supports_reasoning = supports_reasoning, + reasoning_always_on = reasoning_always_on, + _request_reasoning_kwargs = ( + lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None + ), + ), + ) + + def test_split_think_markers_stream_as_reasoning_and_visible_text(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "pla"}}]}, + {"choices": [{"delta": {"content": "n33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "33" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning", + "message", + ] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + assert completed["response"]["output"][1]["content"][0]["text"] == "33" + + def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "show x tags"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert reasoning_deltas == [] + assert "".join(event["delta"] for event in text_deltas) == "show x tags" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == ["message"] + assert completed["response"]["output"][0]["content"][0]["text"] == ( + "show x tags" + ) + + def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "show x tags"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks, supports_reasoning = False) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert reasoning_deltas == [] + assert "".join(event["delta"] for event in text_deltas) == "show x tags" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == ["message"] + assert completed["response"]["output"][0]["content"][0]["text"] == ( + "show x tags" + ) + + def test_reasoning_only_streams_as_visible_message_text(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "plan"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "plan" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning", + "message", + ] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + assert completed["response"]["output"][1]["content"][0]["text"] == "plan" + + def test_structured_reasoning_content_streams_as_reasoning(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"reasoning_content": "plan"}}]}, + {"choices": [{"delta": {"content": "33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "33" + completed = self._payloads(lines, "response.completed")[0] + assert completed["response"]["output"][0]["type"] == "reasoning" + assert completed["response"]["output"][1]["type"] == "message" + + def test_structured_reasoning_content_parts_stream_as_reasoning(self, monkeypatch): + chunks = [ + { + "choices": [ + { + "delta": { + "reasoning_content": { + "content": [ + {"type": "reasoning_text", "text": "plan"}, + {"type": "reasoning_text", "text": " next"}, + ] + } + } + } + ] + }, + {"choices": [{"delta": {"content": "33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan next" + assert "".join(event["delta"] for event in text_deltas) == "33" + assert "reasoning_text" not in "".join(event["delta"] for event in reasoning_deltas) + completed = self._payloads(lines, "response.completed")[0] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan next" + assert completed["response"]["output"][1]["content"][0]["text"] == "33" + + def test_tool_first_stream_closes_items_in_output_index_order(self, monkeypatch): + chunks = [ + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + } + } + ] + }, + {"choices": [{"delta": {"content": "done"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + done_events = self._payloads(lines, "response.output_item.done") + assert [event["output_index"] for event in done_events] == [0, 1] + assert [event["item"]["type"] for event in done_events] == ["function_call", "message"] + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == [ + "function_call", + "message", + ] + def test_requests_usage_and_caps_parallel_tool_calls(self, monkeypatch): import routes.inference as inf_mod @@ -678,6 +1103,15 @@ class TestResponsesStreamAdapter: class TestResponsesOutputFunctionCall: + def test_reasoning_output_item_serialises_full_reasoning_content(self): + item = ResponsesOutputReasoning(content = [{"type": "reasoning_text", "text": "plan"}]) + d = item.model_dump() + assert d["type"] == "reasoning" + assert d["id"].startswith("rs_") + assert d["status"] == "completed" + assert d["summary"] == [] + assert d["content"] == [{"type": "reasoning_text", "text": "plan"}] + def test_direct_construction(self): fc = ResponsesOutputFunctionCall( call_id = "call_1", @@ -778,6 +1212,26 @@ class TestCodexStyleRequestShapes: assert len(req.input) == 3 assert isinstance(req.input[1], ResponsesUnknownInputItem) + def test_emitted_reasoning_item_replay_is_dropped_for_local_chat(self): + payload = ResponsesRequest( + input = [ + {"role": "user", "content": "Hi"}, + { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "content": [{"type": "reasoning_text", "text": "plan"}], + }, + {"role": "assistant", "content": "33"}, + {"role": "user", "content": "Continue"}, + ], + ) + + msgs = _normalise_responses_input(payload) + + assert [m.role for m in msgs] == ["user", "assistant", "user"] + assert all("plan" not in (m.content or "") for m in msgs if isinstance(m.content, str)) + def test_unknown_content_part_type_accepted(self): """Unknown content-part types (e.g. future input_audio) validate as ResponsesUnknownContentPart so the request doesn't 422.""" diff --git a/studio/backend/tests/test_s3_dataset.py b/studio/backend/tests/test_s3_dataset.py new file mode 100644 index 0000000000..f47db565ff --- /dev/null +++ b/studio/backend/tests/test_s3_dataset.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the S3 dataset loader (core.training.s3_dataset). + +boto3 is optional and may be absent in CI, so the S3 client is mocked: a fake +client provides a paginator over a synthetic bucket listing and writes files on +download_file. No network or real AWS credentials are involved. +""" + +import importlib.util +import os +from pathlib import Path + +import pytest + +# Load the modules under test directly by path. Importing them through their +# packages (core.training / models) would execute heavy package __init__ chains +# (structlog, torch, …) that aren't needed for these unit tests. +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load(mod_name, rel_path): + spec = importlib.util.spec_from_file_location(mod_name, _BACKEND / rel_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +s3_dataset = _load("s3_dataset", "core/training/s3_dataset.py") +S3Config = _load("models_training_s3", "models/training.py").S3Config + + +class _FakePaginator: + def __init__(self, keys): + self._keys = keys + + def paginate(self, **kwargs): + prefix = kwargs.get("Prefix") + contents = [{"Key": k} for k in self._keys if prefix is None or k.startswith(prefix)] + # Emit in two pages to exercise pagination handling. + mid = len(contents) // 2 + yield {"Contents": contents[:mid]} + yield {"Contents": contents[mid:]} + + +class _FakeS3Client: + def __init__(self, keys): + self._keys = keys + self.downloaded = [] + + def get_paginator(self, name): + assert name == "list_objects_v2" + return _FakePaginator(self._keys) + + def download_file(self, bucket, key, local_path, **kwargs): + self.downloaded.append((bucket, key, local_path)) + callback = kwargs.get("Callback") + if callback is not None: + callback(1) + with open(local_path, "w", encoding = "utf-8") as f: + f.write(f"content-of:{key}") + + +@pytest.fixture +def fake_client(monkeypatch): + """Force boto3_available True and stub the client builder.""" + keys = [ + "datasets/train.parquet", + "datasets/extra.parquet", + "datasets/notes.txt", # filtered out (unsupported) + "datasets/subdir/", # directory placeholder, skipped + "other/ignore.parquet", # filtered out by prefix + ] + client = _FakeS3Client(keys) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + return client + + +def _cfg(**overrides): + base = { + "bucket": "my-bucket", + "region": "us-east-1", + "prefix": "datasets/", + "access_key_id": "AKIA_TEST", + "secret_access_key": "secret", + "use_iam_role": False, + } + base.update(overrides) + return base + + +def test_downloads_only_supported_files_under_prefix(fake_client, tmp_path): + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + names = sorted(os.path.basename(f) for f in files) + # txt is unsupported, the directory placeholder is skipped, and the + # "other/" key is excluded by the prefix filter. + assert names == ["extra.parquet", "train.parquet"] + for f in files: + assert os.path.exists(f) + + +def test_allows_json_and_jsonl_family(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/train.json", "datasets/extra.jsonl"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert sorted(os.path.basename(f) for f in files) == ["extra.jsonl", "train.json"] + + +def test_ignores_common_json_metadata_files(monkeypatch, tmp_path): + client = _FakeS3Client( + [ + "datasets/train.parquet", + "datasets/schema.json", + "datasets/metadata.json", + "datasets/dataset_info.json", + ] + ) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert [os.path.basename(f) for f in files] == ["train.parquet"] + + +def test_raises_when_prefix_contains_mixed_formats(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/train.parquet", "datasets/stray.csv"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + with pytest.raises(ValueError, match = "mixed dataset formats"): + s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert client.downloaded == [] + + +def test_raises_when_no_supported_files(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/readme.txt"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + with pytest.raises(ValueError, match = "No supported dataset files"): + s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + +def test_raises_when_boto3_missing(monkeypatch, tmp_path): + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: False) + with pytest.raises(RuntimeError, match = "requires boto3"): + s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + +def test_basename_collisions_are_disambiguated(monkeypatch, tmp_path): + # Two keys share a basename under different sub-prefixes. + client = _FakeS3Client(["datasets/a/train.parquet", "datasets/b/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + assert len(files) == 2 + assert len(set(files)) == 2 # no overwrite + + +def test_basename_collision_skips_existing_generated_suffix(monkeypatch, tmp_path): + client = _FakeS3Client( + [ + "datasets/a/train.parquet", + "datasets/b/train_1.parquet", + "datasets/c/train.parquet", + ] + ) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert [os.path.basename(f) for f in files] == [ + "train.parquet", + "train_1.parquet", + "train_2.parquet", + ] + assert len(set(files)) == 3 + assert (tmp_path / "train_1.parquet").read_text(encoding = "utf-8") == ( + "content-of:datasets/b/train_1.parquet" + ) + assert (tmp_path / "train_2.parquet").read_text(encoding = "utf-8") == ( + "content-of:datasets/c/train.parquet" + ) + + +def test_download_handle_cleans_owned_temp_dir(monkeypatch, tmp_path): + target_dir = tmp_path / "owned-download" + client = _FakeS3Client(["datasets/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir)) + + download = s3_dataset.prepare_s3_dataset_download(_cfg()) + + assert target_dir.exists() + assert download.files == [str(target_dir / "train.parquet")] + download.cleanup() + assert not target_dir.exists() + + +def test_dest_dir_is_not_removed_by_cleanup(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + download = s3_dataset.prepare_s3_dataset_download(_cfg(), dest_dir = str(tmp_path)) + + download.cleanup() + assert tmp_path.exists() + assert (tmp_path / "train.parquet").exists() + + +def test_cancel_callback_aborts_and_removes_temp_dir(monkeypatch, tmp_path): + target_dir = tmp_path / "cancelled-download" + client = _FakeS3Client(["datasets/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir)) + calls = 0 + + def cancel_after_download_starts(): + nonlocal calls + calls += 1 + return calls >= 4 + + with pytest.raises(s3_dataset.S3DownloadCancelled): + s3_dataset.prepare_s3_dataset_download( + _cfg(), + cancel_callback = cancel_after_download_starts, + ) + + assert not target_dir.exists() + + +# ── S3Config model (camelCase aliases + credential validation) ── + + +def test_s3config_accepts_camelcase_aliases(): + cfg = S3Config.model_validate( + { + "bucket": "b", + "region": "eu-west-1", + "accessKeyId": "AKIA", + "secretAccessKey": "shh", + } + ) + assert cfg.access_key_id == "AKIA" + assert cfg.secret_access_key == "shh" + # model_dump() yields snake_case for the loader. + assert cfg.model_dump()["access_key_id"] == "AKIA" + + +def test_s3config_accepts_snake_case(): + cfg = S3Config.model_validate( + {"bucket": "b", "access_key_id": "AKIA", "secret_access_key": "shh"} + ) + assert cfg.access_key_id == "AKIA" + + +def test_s3config_requires_credentials_or_iam(): + with pytest.raises(ValueError): + S3Config.model_validate({"bucket": "b"}) + + +def test_s3config_iam_role_needs_no_keys(): + cfg = S3Config.model_validate({"bucket": "b", "useIamRole": True}) + assert cfg.use_iam_role is True diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 8aa6e5df4e..12731783a0 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -28,6 +28,8 @@ from core.inference.tool_call_parser import ( parse_tool_calls_from_text, strip_tool_markup, ) +from state import tool_approvals +from state.tool_approvals import resolve_tool_decision from utils.datasets import is_gpt_oss_model_name @@ -84,8 +86,7 @@ class TestParser: # A code parameter with a literal must not truncate: the # parser uses end-of-body as the only boundary for single-param calls. text = ( - "html = ''\n" - "print('hi')" + "html = ''\nprint('hi')" ) result = parse_tool_calls_from_text(text) assert len(result) == 1 @@ -1033,6 +1034,50 @@ class TestGuardrails: _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "x"})] + def test_confirm_tool_calls_close_after_prompt_cleans_slot(self, monkeypatch): + approval_id = "approval-close-sf" + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: approval_id) + + loop, exec_fn = _make_loop( + turns = [['{"name":"python","arguments":{"code":"print(1)"}}']], + exec_results = ["OK"], + confirm_tool_calls = True, + session_id = "sess", + max_tool_iterations = 1, + ) + + with tool_approvals._lock: + tool_approvals._pending.clear() + + try: + assert next(loop)["type"] == "status" + start = next(loop) + assert start["type"] == "tool_start" + assert start["approval_id"] == approval_id + with tool_approvals._lock: + assert approval_id in tool_approvals._pending + finally: + loop.close() + + with tool_approvals._lock: + assert approval_id not in tool_approvals._pending + assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False + assert exec_fn.calls == [] + + def test_confirm_tool_calls_skips_rag_autoinject(self, monkeypatch): + def fail_autoinject(*_args, **_kwargs): + raise AssertionError("RAG autoinject must not run before approval") + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + loop, exec_fn = _make_loop( + turns = [["plain answer"]], + confirm_tool_calls = True, + rag_scope = {"thread_id": "t1"}, + ) + events = _collect_events(loop) + assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events) + assert exec_fn.calls == [] + def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self): turns = iter( [ diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 27f695b744..928b636e3e 100644 --- a/studio/backend/tests/test_studio_api.py +++ b/studio/backend/tests/test_studio_api.py @@ -149,6 +149,7 @@ def test_help_output(): "--host", "--frontend", "--silent", + "--tensor-parallel", ]: assert flag in out, f"Missing flag {flag!r} in --help output" print(" PASS --help shows all flags") diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py new file mode 100644 index 0000000000..30bfb91a08 --- /dev/null +++ b/studio/backend/tests/test_tensor_parallel.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the Tensor Parallelism toggle. + +The toggle threads a single ``tensor_parallel`` bool from the chat UI +through the load request to a ``--split-mode tensor`` llama-server flag, +and round-trips it back via the load/status responses so the switch +reflects what is actually running. These tests pin: + + * the pydantic request/response/status contract (snake_case key, + default False), + * the backend ``tensor_parallel`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that ``--split-mode tensor`` is emitted only behind the toggle. +""" + +from __future__ import annotations + +import asyncio +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) +_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) +_httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_server_args import resolve_tensor_parallel +from core.inference.tensor_fallback import load_with_tensor_fallback +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default False) ──────────────── + + +def test_load_request_defaults_tensor_parallel_false(): + req = LoadRequest(model_path = "owner/repo") + assert req.tensor_parallel is False + + +def test_load_request_accepts_tensor_parallel(): + req = LoadRequest(model_path = "owner/repo", tensor_parallel = True) + assert req.tensor_parallel is True + + +def test_load_request_round_trips_json_key(): + # The frontend sends the snake_case key verbatim. + req = LoadRequest.model_validate({"model_path": "owner/repo", "tensor_parallel": True}) + assert req.tensor_parallel is True + assert req.model_dump()["tensor_parallel"] is True + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_tensor_parallel(model_cls): + # Default False, and the key is always present in the JSON body. + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + on = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + tensor_parallel = True, + ) + else: + default = model_cls() + on = model_cls(tensor_parallel = True) + assert default.model_dump()["tensor_parallel"] is False + assert on.model_dump()["tensor_parallel"] is True + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_tensor_parallel_property_defaults_false(): + assert LlamaCppBackend().tensor_parallel is False + + +def test_tensor_parallel_property_reflects_field(): + backend = LlamaCppBackend() + backend._tensor_parallel = True + assert backend.tensor_parallel is True + + +def test_unload_resets_tensor_parallel(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._tensor_parallel = True + backend.unload_model() + assert backend.tensor_parallel is False + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(tensor_parallel: bool) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._tensor_parallel = tensor_parallel + return backend + + +def _target_state(backend: LlamaCppBackend, tensor_parallel: bool) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + tensor_parallel = tensor_parallel, + ) + + +@pytest.mark.parametrize("flag", [True, False]) +def test_already_in_target_state_matches_same_tensor_parallel(flag): + assert _target_state(_loaded_backend(flag), flag) is True + + +@pytest.mark.parametrize( + "loaded,requested", + [(False, True), (True, False)], +) +def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, requested): + # Flipping the toggle either direction must force a reload so the + # command is rebuilt with/without --split-mode tensor. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_reconciles_split_mode_extras(): + # Tensor engaged via --split-mode in extras (boolean omitted/default False) + # must match a server already running tensor mode -- no spurious reload. + backend = _loaded_backend(tensor_parallel = True) + backend._extra_args = ["--split-mode", "tensor"] + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = ["--split-mode", "tensor"], + is_vision = False, + tensor_parallel = False, + ) + is True + ) + + +# ── --split-mode tensor is emitted only behind the toggle ──────────── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_split_mode_tensor_is_gated_on_the_toggle(): + src = _load_model_source() + assert ( + 'cmd.extend(["--split-mode", "tensor"])' in src + ), "the tensor-parallel flag emission must be present in load_model" + # The emission lives behind `if tensor_parallel:` -- it must never be + # part of the unconditional base cmd list. + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] if base_end > base_start else "" + assert ( + "--split-mode" not in base_block + ), "--split-mode must be conditional, not in the base cmd list" + gate = src.find("if tensor_parallel:") + emit = src.find('cmd.extend(["--split-mode", "tensor"])') + assert 0 <= gate < emit, "emission must sit under `if tensor_parallel:`" + + +def test_proportional_tensor_split_is_emitted_in_tensor_mode(): + # Asymmetric GPUs (e.g. 48 GB + 24 GB) OOM the smaller card under the + # even default; the allocator weights --tensor-split by free VRAM. Pin + # that the flag is emitted from inside the tensor-parallel block. + src = _load_model_source() + assert '"--tensor-split"' in src + gate = src.find("if tensor_parallel:") + ts = src.find('"--tensor-split"') + nxt_else = src.find("self._tensor_parallel = False") + assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + + +# ── tensor-mode allocation: conservative VRAM budget ───────────────── + + +def _kv_seeded_backend() -> LlamaCppBackend: + # Minimal GGUF metadata so _can_estimate_kv() is True (legacy KV path). + backend = LlamaCppBackend() + backend._n_layers = 32 + backend._embedding_length = 4096 + backend._n_heads = 32 + backend._n_kv_heads = 8 + backend._context_length = 131072 + return backend + + +def test_fit_context_budget_frac_override_is_tighter(): + backend = _kv_seeded_backend() + model_size = 8 * 1024**3 + pool_mib = 24 * 1024 # tight enough that KV capping bites + + fit_default = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") + fit_tp = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16", budget_frac = 0.80) + assert fit_tp < 131072, "expected the context to be capped at this VRAM tier" + assert fit_tp <= fit_default, "a tighter budget must not allow MORE context" + # Omitting the override must reproduce the default budget exactly. + assert backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") == fit_default + + +# ── unsupported-arch load failure -> clean message ─────────────────── + + +def test_split_mode_tensor_arch_failure_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + "llama_model_create: LLAMA_SPLIT_MODE_TENSOR not implemented for " + "architecture 'deepseek2'", + None, + "unsloth/DeepSeek-V3-GGUF", + ) + assert "Tensor parallelism is not supported" in msg + + +def test_unrelated_arch_failure_not_hijacked_by_tensor_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + "unknown model architecture: 'flux'", "/models/flux.gguf", None + ) + assert "Tensor parallelism" not in msg + + +# ── _plan_tensor_parallel: the allocation math (pure, no model/GPU) ─── +# Seeded full-attention KV (~128 KiB/token) via _kv_seeded_backend, so the +# context cap + split are deterministic. Asserts relationships rather than +# magic numbers so the KV estimate can evolve without breaking these. + +_GB = 1024**3 +_ASYM = [(0, 48000), (1, 24000)] # asymmetric pool, 72000 MiB +_SYM = [(0, 24000), (1, 24000)] # symmetric pool + + +def _plan( + model_gb, + target = 131072, + gpus = _ASYM, + mtp = False, +): + b = _kv_seeded_backend() + return b, b._plan_tensor_parallel(gpus, int(model_gb * _GB), target, mtp_engaged = mtp) + + +def _kv_budget_b(model_gb, gpus = _ASYM): + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB) + + +def test_tp_plan_weighted_split_on_asymmetric_big_model(): + b, (ec, mac, gi, ts) = _plan(50) + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + assert gi == [0, 1] + # split weighted by (free - buffer), not raw free + assert ts == [48000 - reserve, 24000 - reserve] + assert ec < 131072 # capped below native + + +def test_tp_plan_even_split_when_model_fits(): + # A small model whose even share fits the smallest GPU -> llama.cpp's even + # default (None), which is safe for archs that crash on a weighted split. + _, (ec, mac, gi, ts) = _plan(4) + assert ts is None + + +def test_tp_plan_symmetric_gpus_use_even_split(): + _, (ec, mac, gi, ts) = _plan(8, gpus = _SYM) + assert ts is None + + +def test_tp_plan_context_fits_pool_budget_no_oom(): + b, (ec, mac, gi, ts) = _plan(50) + # the chosen context's KV must fit the pooled budget (weights + buffers) + assert b._estimate_kv_cache_bytes(ec) <= _kv_budget_b(50) + + +def test_tp_plan_uses_available_vram_not_wasteful(): + # when the cap engages, the chosen context nearly fills the budget + b, (ec, mac, gi, ts) = _plan(50) + assert b._estimate_kv_cache_bytes(ec) >= 0.9 * _kv_budget_b(50) + + +def test_tp_plan_weights_exceed_pool_floors_context(): + # 70 GB > pool minus per-GPU reserves -> floor (triggers layer fallback) + _, (ec, mac, gi, ts) = _plan(70) + assert ec == 2048 + + +def test_tp_plan_floor_never_exceeds_explicit_small_context(): + # An explicit context below the 2048 floor must not be raised: a caller + # asking for 1024 should not have KV sized for 2048 (avoidable OOM). + _, (ec, mac, gi, ts) = _plan(70, target = 1024) # weights exceed pool -> floor path + assert ec == 1024 + _, (ec2, *_rest) = _plan(50, target = 1024) # cap path with a tiny budget + assert ec2 <= 1024 + + +def test_tp_plan_explicit_context_honored_when_it_fits(): + _, (ec, mac, gi, ts) = _plan(50, target = 8192) + assert ec == 8192 + + +def test_tp_plan_explicit_context_capped_when_too_large(): + _, (ec, mac, gi, ts) = _plan(50, target = 131072) + assert 2048 <= ec < 131072 + + +def test_tp_plan_max_available_ctx_reports_native_not_explicit_ctx(): + # An explicit small ctx caps effective_ctx but the UI ceiling + # (max_available_ctx) must reflect the native/hardware cap, not the request. + b = _kv_seeded_backend() + ec, mac, _gi, _ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 8192, max_target_ctx = 131072) + _, native_mac, *_ = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec == 8192 # explicit request honored for the load + assert mac == native_mac > ec # ceiling reflects the hardware cap + + +def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): + _, (ec_no, *_rest) = _plan(50) + _, (ec_mtp, *_rest) = _plan(50, mtp = True) + assert ec_mtp < ec_no + + +def test_tp_plan_no_kv_metadata_floors_context(): + b = LlamaCppBackend() # no KV metadata -> can't size safely + ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec <= 4096 + + +def test_tp_plan_single_gpu_never_splits(): + # The toggle is a no-op without >= 2 GPUs (most dev/CI machines). Even if + # the planner is reached, it must not emit a tensor split. + b = _kv_seeded_backend() + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 24000)], int(8 * _GB), 8192) + assert ts is None + assert gi == [0] + + +def test_tp_plan_zero_gpus_never_splits(): + b = _kv_seeded_backend() + ec, mac, gi, ts = b._plan_tensor_parallel([], int(8 * _GB), 8192) + assert ts is None + assert gi == [] + + +def test_tp_plan_drops_gpu_below_buffer_reserve(): + # A GPU with less free VRAM than the per-device compute-buffer reserve + # can't host tensor mode; it's excluded, which here leaves <2 usable -> no + # split (and gpu_indices reflects only the usable device). + b = _kv_seeded_backend() + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, reserve - 1)], int(8 * _GB), 8192) + assert gi == [0] + assert ts is None + + +# ── route auto-fallback survives a *raised* tensor-load crash ───────── +# A tensor-incompatible model makes load_model RAISE (Gemma 3n aborts) rather +# than return False. The /load fallback helper must catch that and retry with +# layer split -- stripping any --split-mode from the extras so the retry can't +# relaunch tensor -- while a non-tensor load propagates its exception. These +# exercise the real helper with a fake loader (no GPU, no llama-server). + + +class _RecordingLoader: + """Fake ``attempt_load``: crashes whenever tensor mode is effectively + engaged (via the bool or a ``--split-mode`` in extras), like a real + tensor-incompatible model; succeeds on layer split.""" + + def __init__(self): + self.calls: list[tuple] = [] + + async def __call__(self, tensor_parallel, extra_args): + self.calls.append((tensor_parallel, list(extra_args) if extra_args else extra_args)) + if resolve_tensor_parallel(extra_args, tensor_parallel): + raise RuntimeError("llama-server failed to start") + return True + + +def test_tensor_fallback_retries_layer_on_crash(): + loader = _RecordingLoader() + ok = asyncio.run( + load_with_tensor_fallback(loader, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is True + # tensor first (crashes), then layer split. + assert [c[0] for c in loader.calls] == [True, False] + + +def test_tensor_fallback_no_retry_on_success(): + calls: list[bool] = [] + + async def _ok(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return True + + ok = asyncio.run( + load_with_tensor_fallback(_ok, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is True + assert calls == [True] # no fallback when the tensor load succeeds + + +def test_tensor_fallback_retries_when_tensor_returns_false(): + # load_model can signal failure by *returning False* (not only by raising); + # that must trigger the layer-split retry just like a crash does. + calls: list[bool] = [] + + async def _false_on_tensor(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return not resolve_tensor_parallel(extra_args, tensor_parallel) + + ok = asyncio.run( + load_with_tensor_fallback( + _false_on_tensor, requested_tensor = True, extra_args = None, label = "m" + ) + ) + assert ok is True + assert calls == [True, False] + + +def test_tensor_fallback_returns_false_when_both_attempts_fail(): + # Tensor fails and the layer retry also fails -> the helper returns False so + # the route raises its own HTTP 500 (it does not crash mid-flight). + calls: list[bool] = [] + + async def _always_false(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return False + + ok = asyncio.run( + load_with_tensor_fallback(_always_false, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is False + assert calls == [True, False] # tried tensor, then layer split + + +def test_tensor_fallback_skips_layer_retry_when_cancelled(): + # load_model returns False on a user cancellation too. When cancelled() is + # True, the helper must NOT relaunch the load the user just cancelled. + calls: list[bool] = [] + + async def _false_on_tensor(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return False + + ok = asyncio.run( + load_with_tensor_fallback( + _false_on_tensor, + requested_tensor = True, + extra_args = None, + label = "m", + cancelled = lambda: True, + ) + ) + assert ok is False + assert calls == [True] # no layer-split retry after cancellation + + +@pytest.mark.parametrize( + "extras", + [ + ["--split-mode", "tensor", "-c", "4096"], + ["-sm", "tensor", "-c", "4096"], + ["--split-mode=tensor", "-c", "4096"], + ["-sm=tensor", "-c", "4096"], + ], +) +def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras): + # Tensor engaged via extras (boolean False); the retry must drop every + # --split-mode form (long/short, space/=) but keep the user's other flags, + # else resolve_tensor_parallel re-enables tensor and relaunches the crash. + loader = _RecordingLoader() + ok = asyncio.run( + load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m") + ) + assert ok is True + assert len(loader.calls) == 2 + assert loader.calls[1][1] == ["-c", "4096"] # split-mode stripped, -c kept + + +def test_tensor_fallback_propagates_non_tensor_crash(): + async def _always_raise(tensor_parallel, extra_args): + raise RuntimeError("bad model") + + with pytest.raises(RuntimeError, match = "bad model"): + asyncio.run( + load_with_tensor_fallback( + _always_raise, requested_tensor = False, extra_args = None, label = "m" + ) + ) diff --git a/studio/backend/tests/test_tool_approvals.py b/studio/backend/tests/test_tool_approvals.py new file mode 100644 index 0000000000..af792e652c --- /dev/null +++ b/studio/backend/tests/test_tool_approvals.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Concurrency tests for the per-call tool-call confirmation gate. + +``state.tool_approvals`` coordinates two threads: the agentic loop thread +blocked in ``wait_tool_decision`` and the request thread that delivers the +user's choice through ``resolve_tool_decision``. Each gated call carries a +unique ``approval_id`` so a stale or concurrent confirmation can never +resolve the wrong call. These tests exercise that handshake directly -- +no model, no server -- so the race windows are fast and deterministic. +""" + +import threading +import time + +import pytest + +from state import tool_approvals +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + request_tool_decision, + resolve_tool_decision, + wait_tool_decision, +) + + +@pytest.fixture(autouse = True) +def _clear_pending(): + """Each test starts and ends with an empty ``_pending`` map.""" + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +class _Waiter: + """Run ``request_tool_decision`` in a thread and capture its result.""" + + def __init__( + self, + session_id, + approval_id, + cancel_event = None, + timeout = None, + ): + self.session_id = session_id + self.approval_id = approval_id + self.cancel_event = cancel_event + self.timeout = timeout + self.result = None + self._thread = threading.Thread(target = self._run, daemon = True) + + def _run(self): + kwargs = {"cancel_event": self.cancel_event} + if self.timeout is not None: + kwargs["timeout"] = self.timeout + self.result = request_tool_decision(self.session_id, self.approval_id, **kwargs) + + def start(self): + self._thread.start() + _wait_until(lambda: _has_pending(self.approval_id)) + return self + + def join(self, timeout = 5.0): + self._thread.join(timeout = timeout) + assert not self._thread.is_alive(), "waiter thread did not finish" + return self.result + + +def _has_pending(approval_id) -> bool: + with tool_approvals._lock: + return approval_id in tool_approvals._pending + + +def _wait_until( + pred, + timeout = 2.0, + interval = 0.005, +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return True + time.sleep(interval) + return False + + +# ── Basic allow / deny ─────────────────────────────────────────────── + + +def test_allow_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + assert w.join() == "allow" + + +def test_deny_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "deny", session_id = "sess") is True + assert w.join() == "deny" + + +def test_slot_cleaned_up_after_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + resolve_tool_decision(aid, "allow") + w.join() + assert _wait_until(lambda: not _has_pending(aid)) + + +def test_abort_tool_decision_removes_unwaited_slot(): + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + abort_tool_decision(slot, aid) + assert not _has_pending(aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is False + + +def test_approval_ids_are_unique(): + ids = {new_approval_id() for _ in range(1000)} + assert len(ids) == 1000 + + +# ── Pre-registration race (begin before wait) ──────────────────────── + + +def test_resolve_before_wait_is_not_lost(): + """A decision delivered after ``begin`` but before ``wait`` survives. + + The loop registers the slot before it yields ``tool_start``, so even a + confirmation that races ahead of the blocking ``wait`` is recorded on + the slot and returned -- never dropped. + """ + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + # wait() is only entered now, after the decision already landed. + assert wait_tool_decision(slot, aid) == "allow" + assert not _has_pending(aid) + + +# ── Resolver edge cases ────────────────────────────────────────────── + + +def test_resolve_unknown_approval_returns_false(): + assert resolve_tool_decision(new_approval_id(), "allow") is False + + +def test_resolve_empty_approval_returns_false(): + assert resolve_tool_decision("", "allow") is False + assert resolve_tool_decision(None, "allow") is False + + +def test_resolve_wrong_session_scope_returns_false(): + aid = new_approval_id() + w = _Waiter("sess-a", aid).start() + # Correct approval_id but the wrong session must not resolve it. + assert resolve_tool_decision(aid, "allow", session_id = "sess-b") is False + assert _has_pending(aid) + # The right session still works. + assert resolve_tool_decision(aid, "allow", session_id = "sess-a") is True + assert w.join() == "allow" + + +def test_duplicate_resolve_after_completion_returns_false(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "allow") is True + w.join() + assert _wait_until(lambda: not _has_pending(aid)) + assert resolve_tool_decision(aid, "deny") is False + + +def test_first_decision_is_immutable(): + """A second confirmation cannot flip an already-recorded decision. + + The waiter reads ``slot["decision"]`` outside the lock and then cleans up, + so a duplicate or out-of-order POST that lands in that window must be + rejected and must not overwrite the first decision -- an Allow can never + become a Deny. Distinct from the after-completion case above: here the slot + is still pending (no waiter has consumed it yet). + """ + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + # Second decision, same id, before any waiter consumes/cleans the slot. + assert resolve_tool_decision(aid, "deny", session_id = "sess") is False + assert slot["decision"] == "allow" + # The waiter still observes the first (immutable) decision. + assert wait_tool_decision(slot, aid) == "allow" + assert not _has_pending(aid) + + +# ── Cancellation and timeout ───────────────────────────────────────── + + +def test_cancel_event_breaks_wait_as_deny(): + cancel = threading.Event() + aid = new_approval_id() + w = _Waiter("sess", aid, cancel_event = cancel).start() + cancel.set() + assert w.join(timeout = 3.0) == "deny" + assert _wait_until(lambda: not _has_pending(aid)) + + +def test_timeout_returns_deny(): + aid = new_approval_id() + start = time.monotonic() + result = request_tool_decision("sess", aid, timeout = 0.1) + assert result == "deny" + assert time.monotonic() - start < 2.0 + assert not _has_pending(aid) + + +# ── Independence across concurrent calls ───────────────────────────── + + +def test_two_pending_calls_same_session_are_independent(): + """Keying on approval_id, not session, keeps concurrent calls distinct. + + Resolving the first call's id must not unblock or alter the second + call pending in the same session. + """ + a1, a2 = new_approval_id(), new_approval_id() + w1 = _Waiter("sess", a1).start() + w2 = _Waiter("sess", a2).start() + + assert resolve_tool_decision(a1, "deny", session_id = "sess") is True + assert w1.join() == "deny" + # w2 is still waiting on its own id. + assert _has_pending(a2) + assert resolve_tool_decision(a2, "allow", session_id = "sess") is True + assert w2.join() == "allow" + + +def test_concurrent_distinct_calls_route_their_own_decisions(): + n = 25 + waiters = {} + for i in range(n): + aid = new_approval_id() + waiters[aid] = _Waiter(f"s{i}", aid).start() + expected = {aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters)} + for aid, decision in expected.items(): + assert resolve_tool_decision(aid, decision) is True + for aid, w in waiters.items(): + assert w.join() == expected[aid] + + +# ── Constants ──────────────────────────────────────────────────────── + + +def test_rejected_message_is_user_facing_text(): + assert isinstance(TOOL_REJECTED_MESSAGE, str) + assert TOOL_REJECTED_MESSAGE.strip() diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py new file mode 100644 index 0000000000..ce7852c95f --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration tests for the confirmation gate inside the real tool loop. + +These drive ``run_safetensors_tool_loop`` (no model -- hand-crafted fake +generators) with ``confirm_tool_calls=True`` and resolve each pending +decision inline. The slot is registered before ``tool_start`` is yielded, +so resolving right after receiving that event always lands before the +loop blocks. Covers: allow executes once, deny skips execution and feeds +back the rejection, disabled/duplicate calls are not prompted, and a +denied call does not pollute duplicate detection. +""" + +import pytest + +from core.inference.safetensors_agentic import run_safetensors_tool_loop +from state import tool_approvals +from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision + +_SESSION = "loop-session" + + +@pytest.fixture(autouse = True) +def _clear_pending(): + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +class _FakeExecuteTool: + def __init__(self): + self.calls = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + rag_scope = None, + ): + self.calls.append((name, arguments)) + return f"RESULT[{name}]" + + +def _tool_call(name, args_json): + return f'{{"name": "{name}", "arguments": {args_json}}}' + + +def _multi_turn(turns): + """A single_turn generator that yields one full snapshot per turn.""" + turn_iter = iter(turns) + + def _gen(_messages): + try: + yield next(turn_iter) + except StopIteration: + return + + return _gen + + +_DEFAULT_TOOLS = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "web_search"}}, +] + + +def _drive( + turns, + decisions, + *, + tools = None, +): + """Run the loop, resolving each gated tool_start with the next decision. + + The advertised ``tools`` list drives the loop's enabled-tool filter + (pass a list omitting a tool to make a call to it "disabled"). + Returns (events, execute_calls). + """ + decision_iter = iter(decisions) + exec_fn = _FakeExecuteTool() + gen = run_safetensors_tool_loop( + single_turn = _multi_turn(turns), + messages = [{"role": "user", "content": "hi"}], + tools = _DEFAULT_TOOLS if tools is None else tools, + execute_tool = exec_fn, + session_id = _SESSION, + confirm_tool_calls = True, + ) + events = [] + for ev in gen: + events.append(ev) + if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): + # Slot is already registered (begin ran before this yield), so + # the decision lands before the loop enters its blocking wait. + resolve_tool_decision(ev["approval_id"], next(decision_iter), session_id = _SESSION) + return events, exec_fn.calls + + +def _tool_starts(events): + return [e for e in events if e["type"] == "tool_start"] + + +def _tool_ends(events): + return [e for e in events if e["type"] == "tool_end"] + + +def test_allow_executes_the_tool_once(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + ["allow"], + ) + starts = _tool_starts(events) + assert len(starts) == 1 + assert starts[0]["awaiting_confirmation"] is True + assert starts[0]["approval_id"] + assert calls == [("python", {"code": "print(1)"})] + assert _tool_ends(events)[0]["result"] == "RESULT[python]" + + +def test_deny_skips_execution_and_feeds_rejection(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + ["deny"], + ) + assert calls == [] # tool never ran + assert _tool_ends(events)[0]["result"] == TOOL_REJECTED_MESSAGE + + +def test_disabled_tool_is_not_prompted(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + [], + tools = [{"type": "function", "function": {"name": "web_search"}}], + ) + assert _tool_starts(events) == [] + assert _tool_ends(events) == [] + assert calls == [] + + +def test_duplicate_call_is_not_prompted(): + same = _tool_call("python", '{"code": "print(1)"}') + events, calls = _drive([same, same, "final answer"], ["allow"]) + starts = _tool_starts(events) + assert len(starts) == 1 + assert starts[0]["awaiting_confirmation"] is True + assert calls == [("python", {"code": "print(1)"})] + assert len(_tool_ends(events)) == 1 + + +def test_denied_call_can_be_reissued_and_approved(): + # Deny, then the model re-issues the identical call -> approving it must + # execute, not get suppressed as a duplicate (denied calls are not added + # to the duplicate-detection history). + same = _tool_call("python", '{"code": "print(1)"}') + events, calls = _drive([same, same, "final answer"], ["deny", "allow"]) + starts = _tool_starts(events) + assert len(starts) == 2 + assert starts[0]["awaiting_confirmation"] is True + assert starts[1]["awaiting_confirmation"] is True # not treated as dup + assert calls == [("python", {"code": "print(1)"})] # ran once, on approve + ends = _tool_ends(events) + assert ends[0]["result"] == TOOL_REJECTED_MESSAGE + assert ends[1]["result"] == "RESULT[python]" diff --git a/studio/backend/tests/test_tool_confirm_stream.py b/studio/backend/tests/test_tool_confirm_stream.py new file mode 100644 index 0000000000..b8e0472e12 --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_stream.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""End-to-end handshake test for the tool-confirmation gate, no model. + +The real Studio stream wrappers in ``routes/inference.py`` drive the +synchronous agentic generator with ``await asyncio.to_thread(next, gen, +...)`` so the blocking ``threading.Event`` wait runs off the event loop. +This test rebuilds that exact pattern around the real +``state.tool_approvals`` functions, served by a real uvicorn process on +loopback (the same server Studio uses), and proves the load-bearing +property: + +* ``tool_start`` reaches the client before the gate blocks, and +* the separate ``/tool-confirm`` POST is served *while* the stream + connection is blocked, after which the stream resumes with the executed + (allow) or rejected (deny) result -- i.e. no deadlock. + +Each scenario runs under a socket-level timeout, so a regression that +reintroduces a deadlock fails fast instead of hanging the suite. +""" + +import asyncio +import json +import socket +import threading +import time + +import httpx +import pytest +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse + +from state import tool_approvals +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + begin_tool_decision, + new_approval_id, + resolve_tool_decision, + wait_tool_decision, +) + +_EXECUTED_RESULT = "tool executed: 2" + + +@pytest.fixture(autouse = True) +def _clear_pending(): + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +def _build_app() -> FastAPI: + """Minimal app mirroring the real stream/confirm wiring.""" + app = FastAPI() + + def agentic_gen(session_id, cancel_event): + # Same shape as the real loops: register the approval slot, announce + # the call (echoing approval_id), gate on the decision, then either + # execute or feed back the rejection. + approval_id = new_approval_id() + slot = begin_tool_decision(session_id, approval_id) + yield { + "type": "tool_start", + "tool_name": "python", + "approval_id": approval_id, + "awaiting_confirmation": True, + } + denied = wait_tool_decision(slot, approval_id, cancel_event = cancel_event) == "deny" + result = TOOL_REJECTED_MESSAGE if denied else _EXECUTED_RESULT + yield {"type": "tool_end", "tool_name": "python", "result": result} + + @app.post("/stream") + async def stream(req: Request): + body = await req.json() + session_id = body.get("session_id") + cancel_event = threading.Event() + sentinel = object() + + async def wrapper(): + gen = agentic_gen(session_id, cancel_event) + while True: + event = await asyncio.to_thread(next, gen, sentinel) + if event is sentinel: + break + yield f"data: {json.dumps(event)}\n\n" + + return StreamingResponse(wrapper(), media_type = "text/event-stream") + + @app.post("/tool-confirm") + async def tool_confirm(req: Request): + body = await req.json() + resolved = resolve_tool_decision( + body.get("approval_id"), + body.get("decision"), + session_id = body.get("session_id"), + ) + return {"resolved": resolved} + + return app + + +def _free_port() -> int: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +class _Server: + """Run a uvicorn server in a background thread for the test's lifetime.""" + + def __init__(self, app): + self.port = _free_port() + config = uvicorn.Config(app, host = "127.0.0.1", port = self.port, log_level = "warning") + self.server = uvicorn.Server(config) + self._thread = threading.Thread(target = self.server.run, daemon = True) + + def __enter__(self): + self._thread.start() + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + if self.server.started: + return self + time.sleep(0.02) + raise AssertionError("uvicorn did not start in time") + + def __exit__(self, *exc): + self.server.should_exit = True + self._thread.join(timeout = 10.0) + + @property + def base_url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +async def _gate_is_blocking(approval_id) -> None: + """Wait until the stream thread is parked on this approval's slot. + + The slot is registered before ``tool_start`` is yielded, so it exists + by the time the client receives the event -- exactly as in reality, + where the confirm POST only arrives after the card renders. + """ + for _ in range(400): + with tool_approvals._lock: + slot = tool_approvals._pending.get(approval_id) + if slot is not None and not slot["event"].is_set(): + return + await asyncio.sleep(0.005) + raise AssertionError("gate never started waiting") + + +async def _drive(base_url, session_id, decision): + events = [] + resolved = None + timeout = httpx.Timeout(10.0) + async with httpx.AsyncClient(base_url = base_url, timeout = timeout) as client: + async with client.stream("POST", "/stream", json = {"session_id": session_id}) as resp: + assert resp.status_code == 200 + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + event = json.loads(line[len("data: ") :]) + events.append(event) + if event["type"] == "tool_start": + # The stream is now blocked on the gate; the confirm + # POST (echoing approval_id) must still be served over a + # second connection. + approval_id = event["approval_id"] + await _gate_is_blocking(approval_id) + r = await client.post( + "/tool-confirm", + json = { + "session_id": session_id, + "approval_id": approval_id, + "decision": decision, + }, + ) + resolved = r.json()["resolved"] + return events, resolved + + +def _run(session_id, decision): + with _Server(_build_app()) as srv: + return asyncio.run( + asyncio.wait_for(_drive(srv.base_url, session_id, decision), timeout = 15.0) + ) + + +def _types(events): + return [e["type"] for e in events] + + +def test_allow_resumes_stream_with_executed_result(): + events, resolved = _run("sess-allow", "allow") + assert resolved is True + assert _types(events) == ["tool_start", "tool_end"] + assert events[-1]["result"] == _EXECUTED_RESULT + + +def test_deny_resumes_stream_with_rejection_result(): + events, resolved = _run("sess-deny", "deny") + assert resolved is True + assert _types(events) == ["tool_start", "tool_end"] + assert events[-1]["result"] == TOOL_REJECTED_MESSAGE + + +def test_tool_start_precedes_the_block_and_carries_approval_id(): + # The first streamed event is always tool_start, proving the buttons + # can render before the backend pauses for the decision -- and it + # carries the approval_id / awaiting_confirmation the UI needs. + events, _ = _run("sess-order", "allow") + assert events[0]["type"] == "tool_start" + assert events[0]["awaiting_confirmation"] is True + assert events[0]["approval_id"] diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py new file mode 100644 index 0000000000..91fdac9961 --- /dev/null +++ b/studio/backend/tests/test_training_resume.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for resumable training run eligibility.""" + +import importlib.util +import json +from pathlib import Path + + +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load_resume_module(): + spec = importlib.util.spec_from_file_location( + "training_resume_under_test", + _BACKEND / "core" / "training" / "resume.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +resume = _load_resume_module() + + +def _stopped_run(**overrides): + run = { + "status": "stopped", + "final_step": 5, + "total_steps": 10, + "output_dir": "/tmp/unsloth-output", + "resumed_later": False, + "config_json": json.dumps({"hf_dataset": "org/dataset"}), + } + run.update(overrides) + return run + + +def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run()) is True + + +def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run( + config_json = json.dumps( + { + "dataset_source": "s3", + "s3_dataset": { + "bucket": "training-data", + "prefix": "datasets/", + "region": "us-east-1", + "use_iam_role": True, + }, + } + ) + ) + + assert resume.can_resume_run(run) is False + + +def test_can_resume_run_rejects_s3_metadata_marker(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(config_json = json.dumps({"s3_dataset": {"bucket": "training-data"}})) + + assert resume.can_resume_run(run) is False + + +def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + config_json = json.dumps({"dataset_source": "s3", "s3_dataset": {"bucket": "training-data"}}) + + studio_db.create_run( + id = "run-s3", + model_name = "unsloth/test-model", + dataset_name = "s3://training-data", + config_json = config_json, + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + + result = studio_db.list_runs() + + assert result["runs"][0]["config_json"] == config_json diff --git a/studio/backend/utils/hardware/apple.py b/studio/backend/utils/hardware/apple.py new file mode 100644 index 0000000000..62dbd10b8d --- /dev/null +++ b/studio/backend/utils/hardware/apple.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Apple Silicon GPU temperature and power -- no sudo required. + +Mirrors macmon's approach (https://github.com/vladkens/macmon): + * Temperature: average of the AppleSMC "Tg*" float keys (available since + macOS 14; on older systems the keys are absent and this returns None). + * Power: IOReport "Energy Model" group, "GPU Energy" channel. Each poll + diffs the energy counter against the previous poll's sample, so the + result is the average wattage over the polling window. The first poll + only sets the baseline and returns None. + +Public API (never raises; returns None when sensors are unavailable): + read_gpu_temperature_c() + read_gpu_power_w() +""" + +import ctypes +import struct +import time +from typing import Iterable, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +_IOKIT_PATH = "/System/Library/Frameworks/IOKit.framework/IOKit" +_CF_PATH = "/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation" +_IOREPORT_PATH = "/usr/lib/libIOReport.dylib" + +# AppleSMC user-client protocol (same constants as macmon / SMCKit). +_SMC_SELECTOR_HANDLE_EVENT = 2 +_SMC_CMD_READ_BYTES = 5 +_SMC_CMD_KEY_AT_INDEX = 8 +_SMC_CMD_KEY_INFO = 9 + +_MAX_VALID_TEMP_C = 150.0 +_CF_STRING_ENCODING_UTF8 = 0x08000100 +_ENERGY_UNIT_DIVISORS = {"mJ": 1e3, "uJ": 1e6, "nJ": 1e9} + + +# ========== Pure helpers ========== + + +def _fourcc(key: str) -> int: + """Encode a 4-char SMC key/type name as a big-endian integer.""" + return int.from_bytes(key.encode("ascii"), "big") + + +def _fourcc_str(value: int) -> str: + return value.to_bytes(4, "big").decode("ascii", errors = "replace") + + +def _watts(energy: int, unit: str, elapsed_s: float) -> Optional[float]: + """Convert an IOReport energy counter delta into average watts.""" + divisor = _ENERGY_UNIT_DIVISORS.get(unit.strip()) + if divisor is None or elapsed_s <= 0: + return None + return energy / divisor / elapsed_s + + +def _average_valid_temps(values: Iterable[float]) -> Optional[float]: + valid = [v for v in values if 0.0 < v <= _MAX_VALID_TEMP_C] + if not valid: + return None + return round(sum(valid) / len(valid), 1) + + +def _is_gpu_energy_channel(name: str) -> bool: + # Exact "GPU Energy" plus "DIE_N_GPU Energy" on Ultra chips; the separate + # "GPU SRAM*" channels are not GPU core power. + return name.endswith("GPU Energy") and "SRAM" not in name + + +# ========== AppleSMC structs (layout must match the kernel exactly) ========== + + +class _SMCKeyDataVers(ctypes.Structure): + _fields_ = [ + ("major", ctypes.c_uint8), + ("minor", ctypes.c_uint8), + ("build", ctypes.c_uint8), + ("reserved", ctypes.c_uint8), + ("release", ctypes.c_uint16), + ] + + +class _SMCPLimitData(ctypes.Structure): + _fields_ = [ + ("version", ctypes.c_uint16), + ("length", ctypes.c_uint16), + ("cpu_p_limit", ctypes.c_uint32), + ("gpu_p_limit", ctypes.c_uint32), + ("mem_p_limit", ctypes.c_uint32), + ] + + +class _SMCKeyInfo(ctypes.Structure): + _fields_ = [ + ("data_size", ctypes.c_uint32), + ("data_type", ctypes.c_uint32), + ("data_attributes", ctypes.c_uint8), + ] + + +class _SMCKeyData(ctypes.Structure): + _fields_ = [ + ("key", ctypes.c_uint32), + ("vers", _SMCKeyDataVers), + ("p_limit_data", _SMCPLimitData), + ("key_info", _SMCKeyInfo), + ("result", ctypes.c_uint8), + ("status", ctypes.c_uint8), + ("data8", ctypes.c_uint8), + ("data32", ctypes.c_uint32), + ("bytes", ctypes.c_uint8 * 32), + ] + + +# ========== Library loaders ========== + + +def _load_iokit() -> ctypes.CDLL: + iokit = ctypes.CDLL(_IOKIT_PATH) + iokit.IOServiceMatching.restype = ctypes.c_void_p + iokit.IOServiceMatching.argtypes = [ctypes.c_char_p] + iokit.IOServiceGetMatchingServices.argtypes = [ + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint32), + ] + iokit.IOIteratorNext.restype = ctypes.c_uint32 + iokit.IOIteratorNext.argtypes = [ctypes.c_uint32] + iokit.IORegistryEntryGetName.argtypes = [ctypes.c_uint32, ctypes.c_char_p] + iokit.IOServiceOpen.argtypes = [ + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.POINTER(ctypes.c_uint32), + ] + iokit.IOObjectRelease.argtypes = [ctypes.c_uint32] + iokit.IOConnectCallStructMethod.argtypes = [ + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_size_t), + ] + return iokit + + +def _load_cf() -> ctypes.CDLL: + cf = ctypes.CDLL(_CF_PATH) + cf.CFStringCreateWithCString.restype = ctypes.c_void_p + cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint32] + cf.CFStringGetCString.restype = ctypes.c_bool + cf.CFStringGetCString.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_long, + ctypes.c_uint32, + ] + cf.CFRelease.argtypes = [ctypes.c_void_p] + cf.CFDictionaryGetValue.restype = ctypes.c_void_p + cf.CFDictionaryGetValue.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + cf.CFArrayGetCount.restype = ctypes.c_long + cf.CFArrayGetCount.argtypes = [ctypes.c_void_p] + cf.CFArrayGetValueAtIndex.restype = ctypes.c_void_p + cf.CFArrayGetValueAtIndex.argtypes = [ctypes.c_void_p, ctypes.c_long] + return cf + + +def _load_ioreport() -> ctypes.CDLL: + ior = ctypes.CDLL(_IOREPORT_PATH) + ior.IOReportCopyChannelsInGroup.restype = ctypes.c_void_p + ior.IOReportCopyChannelsInGroup.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint64, + ctypes.c_uint64, + ctypes.c_uint64, + ] + ior.IOReportCreateSubscription.restype = ctypes.c_void_p + ior.IOReportCreateSubscription.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_uint64, + ctypes.c_void_p, + ] + ior.IOReportCreateSamples.restype = ctypes.c_void_p + ior.IOReportCreateSamples.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + ior.IOReportCreateSamplesDelta.restype = ctypes.c_void_p + ior.IOReportCreateSamplesDelta.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + ior.IOReportChannelGetChannelName.restype = ctypes.c_void_p + ior.IOReportChannelGetChannelName.argtypes = [ctypes.c_void_p] + ior.IOReportChannelGetUnitLabel.restype = ctypes.c_void_p + ior.IOReportChannelGetUnitLabel.argtypes = [ctypes.c_void_p] + ior.IOReportSimpleGetIntegerValue.restype = ctypes.c_int64 + ior.IOReportSimpleGetIntegerValue.argtypes = [ctypes.c_void_p, ctypes.c_int32] + return ior + + +def _cfstr(cf: ctypes.CDLL, text: str) -> int: + return cf.CFStringCreateWithCString(None, text.encode("utf-8"), _CF_STRING_ENCODING_UTF8) + + +def _from_cfstr(cf: ctypes.CDLL, ref: Optional[int]) -> str: + if not ref: + return "" + buf = ctypes.create_string_buffer(128) + if not cf.CFStringGetCString(ref, buf, len(buf), _CF_STRING_ENCODING_UTF8): + return "" + return buf.value.decode("utf-8", errors = "replace").strip() + + +# ========== SMC connection (GPU temperature) ========== + + +class _SMCConnection: + """Connection to AppleSMCKeysEndpoint; discovers "Tg*" GPU temp keys once.""" + + def __init__(self): + self._iokit = _load_iokit() + self._conn = self._open() + self._key_info_cache: dict[int, _SMCKeyInfo] = {} + self.gpu_keys = [ + key + for key in self._all_keys() + if key.startswith("Tg") and self.read_float(key) is not None + ] + + def _open(self) -> int: + iterator = ctypes.c_uint32(0) + matching = self._iokit.IOServiceMatching(b"AppleSMC") + if self._iokit.IOServiceGetMatchingServices(0, matching, ctypes.byref(iterator)) != 0: + raise OSError("AppleSMC service not found") + try: + conn = self._open_keys_endpoint(iterator.value) + finally: + self._iokit.IOObjectRelease(iterator.value) + if conn is None: + raise OSError("AppleSMCKeysEndpoint not found") + return conn + + def _open_keys_endpoint(self, iterator: int) -> Optional[int]: + task = ctypes.CDLL(None).mach_task_self() + while device := self._iokit.IOIteratorNext(iterator): + name = ctypes.create_string_buffer(128) + self._iokit.IORegistryEntryGetName(device, name) + if name.value != b"AppleSMCKeysEndpoint": + self._iokit.IOObjectRelease(device) + continue + conn = ctypes.c_uint32(0) + status = self._iokit.IOServiceOpen(device, task, 0, ctypes.byref(conn)) + self._iokit.IOObjectRelease(device) + if status != 0: + raise OSError(f"IOServiceOpen(AppleSMCKeysEndpoint) failed: {status}") + return conn.value + return None + + def _call(self, ival: _SMCKeyData) -> _SMCKeyData: + oval = _SMCKeyData() + olen = ctypes.c_size_t(ctypes.sizeof(_SMCKeyData)) + status = self._iokit.IOConnectCallStructMethod( + self._conn, + _SMC_SELECTOR_HANDLE_EVENT, + ctypes.byref(ival), + ctypes.sizeof(_SMCKeyData), + ctypes.byref(oval), + ctypes.byref(olen), + ) + if status != 0: + raise OSError(f"IOConnectCallStructMethod failed: {status}") + if oval.result != 0: + raise OSError(f"SMC result code: {oval.result}") + return oval + + def _read_key_info(self, key_id: int) -> _SMCKeyInfo: + cached = self._key_info_cache.get(key_id) + if cached is not None: + return cached + oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_KEY_INFO)) + self._key_info_cache[key_id] = oval.key_info + return oval.key_info + + def _read_bytes(self, key: str) -> Optional[bytes]: + try: + key_id = _fourcc(key) + info = self._read_key_info(key_id) + oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_READ_BYTES, key_info = info)) + return bytes(oval.bytes[: info.data_size]) + except OSError: + return None + + def read_float(self, key: str) -> Optional[float]: + try: + info = self._read_key_info(_fourcc(key)) + except OSError: + return None + if info.data_size != 4 or info.data_type != _fourcc("flt "): + return None + data = self._read_bytes(key) + if data is None or len(data) != 4: + return None + return struct.unpack(" Optional[str]: + try: + oval = self._call(_SMCKeyData(data8 = _SMC_CMD_KEY_AT_INDEX, data32 = index)) + return oval.key.to_bytes(4, "big").decode("ascii") + except (OSError, UnicodeDecodeError): + return None + + def _all_keys(self) -> list[str]: + count_bytes = self._read_bytes("#KEY") + if count_bytes is None or len(count_bytes) != 4: + return [] + count = int.from_bytes(count_bytes, "big") + names = (self._key_name_at(i) for i in range(count)) + return [name for name in names if name is not None] + + def gpu_temperature_c(self) -> Optional[float]: + readings = (self.read_float(key) for key in self.gpu_keys) + return _average_valid_temps(value for value in readings if value is not None) + + +# ========== IOReport subscription (GPU power) ========== + + +class _IOReportEnergy: + """Persistent subscription to the "Energy Model" group for GPU wattage.""" + + def __init__(self): + self._cf = _load_cf() + self._ior = _load_ioreport() + self._channels = self._ior.IOReportCopyChannelsInGroup( + _cfstr(self._cf, "Energy Model"), None, 0, 0, 0 + ) + if not self._channels: + raise OSError("IOReport 'Energy Model' channel group unavailable") + subscribed = ctypes.c_void_p() + self._sub = self._ior.IOReportCreateSubscription( + None, self._channels, ctypes.byref(subscribed), 0, None + ) + if not self._sub: + raise OSError("IOReportCreateSubscription failed") + # Sample with the channels IOReport subscribes us to, not the requested + # group (matches macmon); fall back if the OS leaves it unset. + self._sample_channels = subscribed if subscribed else self._channels + self._channels_key = _cfstr(self._cf, "IOReportChannels") + self._prev: Optional[tuple[int, float]] = None # (sample ref, monotonic s) + + def gpu_power_w(self) -> Optional[float]: + sample = self._ior.IOReportCreateSamples(self._sub, self._sample_channels, None) + if not sample: + return None + now = time.monotonic() + prev, self._prev = self._prev, (sample, now) + if prev is None: + return None + prev_sample, prev_time = prev + delta = self._ior.IOReportCreateSamplesDelta(prev_sample, sample, None) + self._cf.CFRelease(prev_sample) + if not delta: + return None + try: + return self._gpu_watts_from_delta(delta, now - prev_time) + finally: + self._cf.CFRelease(delta) + + def _gpu_watts_from_delta(self, delta: int, elapsed_s: float) -> Optional[float]: + items = self._cf.CFDictionaryGetValue(delta, self._channels_key) + if not items: + return None + total: Optional[float] = None + for i in range(self._cf.CFArrayGetCount(items)): + item = self._cf.CFArrayGetValueAtIndex(items, i) + name = _from_cfstr(self._cf, self._ior.IOReportChannelGetChannelName(item)) + if not _is_gpu_energy_channel(name): + continue + unit = _from_cfstr(self._cf, self._ior.IOReportChannelGetUnitLabel(item)) + energy = self._ior.IOReportSimpleGetIntegerValue(item, 0) + watts = _watts(energy, unit, elapsed_s) + if watts is not None: + total = (total or 0.0) + watts + if total is None or total < 0: # negative = counter reset; show -- not a bogus draw + return None + return round(total, 1) + + +# ========== Public API (module singletons, failure-latched) ========== + +_smc: Optional[_SMCConnection] = None +_smc_failed = False +_energy: Optional[_IOReportEnergy] = None +_energy_failed = False + + +def read_gpu_temperature_c() -> Optional[float]: + """Average Apple GPU die temperature in degrees C, or None if unavailable.""" + global _smc, _smc_failed + if _smc_failed: + return None + try: + if _smc is None: + _smc = _SMCConnection() + return _smc.gpu_temperature_c() + except Exception as e: + _smc_failed = True + logger.warning("Apple SMC GPU temperature unavailable: %s", e) + return None + + +def read_gpu_power_w() -> Optional[float]: + """Average GPU power in watts since the previous call, or None. + + The first call establishes the baseline sample and returns None. + """ + global _energy, _energy_failed + if _energy_failed: + return None + try: + if _energy is None: + _energy = _IOReportEnergy() + return _energy.gpu_power_w() + except Exception as e: + _energy_failed = True + logger.warning("Apple IOReport GPU power unavailable: %s", e) + return None diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 893d0364f7..86dfa8a93f 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -714,17 +714,19 @@ def get_gpu_utilization() -> Dict[str, Any]: except Exception: pass + from . import apple + return { "available": True, "backend": device.value, "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, - "temperature_c": None, + "temperature_c": apple.read_gpu_temperature_c(), "vram_used_gb": round(vram_used_gb, 2), "vram_total_gb": round(total_gb, 2), "vram_utilization_pct": ( round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None ), - "power_draw_w": None, + "power_draw_w": apple.read_gpu_power_w(), "power_limit_w": None, "power_utilization_pct": None, } diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index 3e5066ca2d..87d0d2ec01 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import os +import re import time from datetime import datetime, timezone from pathlib import Path @@ -104,11 +105,18 @@ def _save_disk_cache(repo: str, latest_tag: Optional[str]) -> None: def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]: - """GitHub API call. None on any failure (offline, rate-limited, etc).""" + """Newest published release tag for `repo`, by publish time. + + Resolves "latest" the way install_llama_prebuilt.py does (newest + non-draft/non-prerelease by ``published_at``), NOT via GitHub's + ``/releases/latest`` pointer. That pointer sorts by commit date and can lag + behind the build the installer actually installs, so detection and apply + disagreed -- the cause of the downgrade/sticky banner. None on any failure + (offline, rate-limited, etc).""" import urllib.error import urllib.request - url = f"https://api.github.com/repos/{repo}/releases/latest" + url = f"https://api.github.com/repos/{repo}/releases?per_page=30" headers = { "Accept": "application/vnd.github+json", "User-Agent": "unsloth-studio-freshness-check", @@ -128,8 +136,21 @@ def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]: ) as exc: logger.debug("freshness fetch failed", repo = repo, error = str(exc)) return None - tag = data.get("tag_name") - return tag if isinstance(tag, str) and tag else None + if not isinstance(data, list): + return None + published = [ + r + for r in data + if isinstance(r, dict) + and not r.get("draft") + and not r.get("prerelease") + and isinstance(r.get("tag_name"), str) + and r.get("tag_name") + ] + if not published: + return None + newest = max(published, key = lambda r: r.get("published_at") or "") + return newest["tag_name"] def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]: @@ -172,19 +193,58 @@ def _parse_installed_at(value: object) -> Optional[datetime]: return dt +def parse_base_build(tag: object) -> Optional[int]: + """Numeric base build from a release tag. Handles both a plain ``bNNNN`` and + a mix-build tag like ``b9596-mix-`` (anchored at the start, so the mix + suffix doesn't defeat it). None for anything not starting with ``bNNNN``.""" + if not isinstance(tag, str): + return None + m = re.match(r"b(\d+)", tag.strip()) + return int(m.group(1)) if m else None + + +def is_behind(installed: Optional[str], latest: Optional[str]) -> bool: + """Whether `installed` is genuinely behind `latest`, comparing the FULL + release identity (so a mix build can legitimately be the latest) with a + base-build guard so a lagging GitHub /releases/latest can never read as an + update or a downgrade. + + - identical tags -> not behind (clears the sticky banner post-update) + - higher base build on `latest` -> behind; lower -> NOT behind (downgrade guard) + - same base build: a different/new mix -> behind, but a bare ``bNNNN`` never + supersedes a mix build (extra PRs) at that base -> not behind + - non-bNNNN tags -> behind (plain inequality, since they already differ) + """ + if not installed or not latest: + return False + installed, latest = installed.strip(), latest.strip() + if installed == latest: + return False + ib, lb = parse_base_build(installed), parse_base_build(latest) + if ib is None or lb is None: + return True + if lb != ib: + return lb > ib + # Same base build, different tags: offer a mix (latest carries a suffix), but + # never offer a bare base over a mix install at the same base. + return latest != f"b{lb}" + + def check_prebuilt_freshness( binary_path: Optional[str], *, threshold_days: int = STALENESS_THRESHOLD_DAYS, now: Optional[datetime] = None, ) -> dict: - """Returns {has_marker, stale, installed_tag, latest_tag, + """Returns {has_marker, stale, behind, installed_tag, latest_tag, installed_at_utc, age_days, published_repo, threshold_days}. - stale = True iff installed != latest AND age >= threshold. - Fails open on missing data (stale stays False).""" + behind = installed genuinely older than latest (see is_behind). + stale = behind AND age >= threshold. + Fails open on missing data (behind/stale stay False).""" out: dict = { "has_marker": False, "stale": False, + "behind": False, "installed_tag": None, "latest_tag": None, "installed_at_utc": None, @@ -196,16 +256,25 @@ def check_prebuilt_freshness( if not marker: return out out["has_marker"] = True + # Display prefers the normalized base ("tag"); comparison below prefers the + # full "release_tag" -- deliberately opposite fallbacks. out["installed_tag"] = marker.get("tag") or marker.get("release_tag") out["installed_at_utc"] = marker.get("installed_at_utc") out["published_repo"] = marker.get("published_repo") + # The marker records both a normalized base tag ("tag", e.g. b9596) and the + # full release tag ("release_tag", e.g. b9596-mix-). Compare against the + # FULL identity, since GitHub /releases/latest returns the full tag_name -- + # comparing the normalized base against the full latest is what produced the + # permanent "downgrade" banner on every mix release. + installed_full = marker.get("release_tag") or marker.get("tag") repo = out["published_repo"] - if not repo or not out["installed_tag"]: + if not repo or not installed_full: return out latest = latest_published_release(repo) out["latest_tag"] = latest - if not latest or latest == out["installed_tag"]: + out["behind"] = is_behind(installed_full, latest) + if not out["behind"]: return out installed_at = _parse_installed_at(out["installed_at_utc"]) @@ -232,7 +301,22 @@ def format_stale_warning(info: dict) -> str: ) -def reset_caches() -> None: - """Test-only: drop all in-memory caches.""" +def reset_caches(*, drop_disk: bool = False) -> None: + """Drop the in-memory freshness caches. The no-arg form is test-only. + + With ``drop_disk = True`` also delete the on-disk 24h release cache. Used by + the post-install/update path: in-memory clearing alone leaves the stale + same-base value on disk, so if the post-install GitHub refresh can't reach + the network, ``latest_published_release`` would replay that stale disk value + (see its last-good fallback) and the banner could linger. Dropping the disk + cache makes latest read as None in that offline case, so the banner fails + open (off) instead of pointing at the just-replaced build.""" _marker_cache.clear() _release_memo.clear() + if drop_disk: + import shutil + + # _cache_dir() is a dedicated freshness-only subdir; it is re-created on + # the next _save_disk_cache. ignore_errors so a missing/locked dir is a + # no-op rather than breaking an otherwise successful install. + shutil.rmtree(_cache_dir(), ignore_errors = True) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index b138b29af3..6eb34ffe34 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -59,10 +59,17 @@ _job: dict = { "from_tag": None, "to_tag": None, "error": None, + "progress": None, "started_at": None, "finished_at": None, } +# Matches the installer's download progress lines, e.g. +# "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s". +_PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(") +# The download dominates the update; extract/validate fill the last slice. +_DOWNLOAD_PROGRESS_CEILING = 0.95 + def _utcnow() -> str: return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) @@ -279,9 +286,10 @@ def get_update_status(*, force_refresh: bool = False) -> dict: freshness = check_prebuilt_freshness(binary) installed = freshness.get("installed_tag") latest = freshness.get("latest_tag") - update_available = bool( - freshness.get("has_marker") and installed and latest and installed != latest - ) + # `behind` compares the full release identity with a base-build guard, so a + # lagging /releases/latest or a mix-tagged latest can't show a false update + # (see llama_cpp_freshness.is_behind). + update_available = bool(freshness.get("has_marker") and freshness.get("behind")) with _job_lock: job = dict(_job) @@ -357,19 +365,58 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path ] cmd.extend(_rocm_install_args(asset)) logger.info("llama update: installing", cmd = " ".join(cmd)) - proc = subprocess.run( + # Stream the installer output so download percent lines feed + # job["progress"]; finer milestones via UNSLOTH_PROGRESS_PERCENT_STEP. + env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + proc = subprocess.Popen( cmd, - capture_output = True, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, text = True, - timeout = _INSTALL_TIMEOUT_SECONDS, + env = env, ) - if proc.returncode != 0: - tail = (proc.stderr or proc.stdout or "").strip()[-1500:] - raise RuntimeError(f"installer exited {proc.returncode}: {tail or 'no output'}") + timed_out = threading.Event() - # New UNSLOTH_PREBUILT_INFO.json is on disk; drop caches so the next - # status read reflects the freshly installed tag. - reset_caches() + def _kill_on_timeout() -> None: + timed_out.set() + proc.kill() + + watchdog = threading.Timer(_INSTALL_TIMEOUT_SECONDS, _kill_on_timeout) + watchdog.daemon = True + watchdog.start() + tail_lines: list[str] = [] + try: + assert proc.stdout is not None + for line in proc.stdout: + tail_lines.append(line) + if len(tail_lines) > 80: + del tail_lines[0] + m = _PROGRESS_LINE_RE.search(line) + if m is None: + continue + fraction = min(float(m.group(1)) / 100.0, 1.0) * _DOWNLOAD_PROGRESS_CEILING + with _job_lock: + _job["progress"] = max(_job.get("progress") or 0.0, fraction) + returncode = proc.wait() + finally: + watchdog.cancel() + if timed_out.is_set(): + raise RuntimeError(f"installer timed out after {_INSTALL_TIMEOUT_SECONDS}s") + if returncode != 0: + tail = "".join(tail_lines).strip()[-1500:] + raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}") + + # New UNSLOTH_PREBUILT_INFO.json is on disk; drop the in-memory AND the + # on-disk freshness caches, then re-prime the 24h disk cache with the + # true newest, so the banner can't linger on a stale same-base value + # after the swap. drop_disk matters when the refresh below can't reach + # GitHub: without it, latest_published_release would replay the stale + # disk value; with it, latest reads as None and the banner fails open. + reset_caches(drop_disk = True) + try: + latest_published_release(repo, force_refresh = True) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: post-install freshness refresh failed", error = str(exc)) new_marker = read_install_marker(_find_binary()) new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag") @@ -382,6 +429,7 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path ), to_tag = new_tag, error = None, + progress = 1.0, finished_at = _utcnow(), ) logger.info("llama update: success", to_tag = new_tag) @@ -417,7 +465,24 @@ def start_update() -> dict: "job": get_update_status()["job"], } + # A job already in flight wins over any freshness re-check below (and skips + # its network call). The final lock block re-checks to close the TOCTOU. + with _job_lock: + if _job["state"] == _JOB_RUNNING: + return {"started": False, "reason": "already_running", "job": dict(_job)} + if marker: + # Mirror the detection guard: a direct POST or a stale banner must not + # start an install when the latest is not actually newer (force a fresh + # check so a stale 24h cache can't wrongly block a real update either). + status = get_update_status(force_refresh = True) + if not status.get("update_available"): + return { + "started": False, + "reason": "up_to_date", + "message": "The installed llama.cpp build is already at the latest prebuilt.", + "job": status["job"], + } install_dir = _install_dir_for(binary) repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO from_tag = marker.get("tag") or marker.get("release_tag") @@ -467,6 +532,7 @@ def start_update() -> dict: from_tag = from_tag, to_tag = None, error = None, + progress = 0.0, started_at = _utcnow(), finished_at = None, ) @@ -491,6 +557,7 @@ def _reset_job_for_tests() -> None: from_tag = None, to_tag = None, error = None, + progress = None, started_at = None, finished_at = None, ) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f61c210cf6..d87fb6aa09 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1371,12 +1371,11 @@ def _iter_hf_cache_snapshots(repo_id: str): return cache_dir = Path(hf_constants.HF_HUB_CACHE) - if not cache_dir.is_dir(): - return - target = f"models--{repo_id.replace('/', '--')}".lower() repo_dir: Optional[Path] = None try: + if not cache_dir.is_dir(): + return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: repo_dir = entry @@ -1387,10 +1386,9 @@ def _iter_hf_cache_snapshots(repo_id: str): return snapshots = repo_dir / "snapshots" - if not snapshots.is_dir(): - return - try: + if not snapshots.is_dir(): + return snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] except OSError: return diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 6913a3be73..eff9b64678 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -43,13 +43,14 @@ from .storage_roots import ( resolve_under_root, resolve_output_dir, resolve_export_dir, + resolve_export_write_dir, resolve_tensorboard_dir, resolve_dataset_path, ) # Re-export shim: mark project-path helpers as used so the import-hoist # safety net does not flag them as unused. -_REEXPORTED = (documents_root, project_workspaces_root) +_REEXPORTED = (documents_root, project_workspaces_root, resolve_export_write_dir) __all__ = [ "normalize_path", @@ -89,6 +90,7 @@ __all__ = [ "resolve_under_root", "resolve_output_dir", "resolve_export_dir", + "resolve_export_write_dir", "resolve_tensorboard_dir", "resolve_dataset_path", ] diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 2d4f5ac243..d336bc2e71 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -6,7 +6,7 @@ from __future__ import annotations import json import os import sys -from pathlib import Path +from pathlib import Path, PurePosixPath, PureWindowsPath import tempfile @@ -317,6 +317,26 @@ def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ( return Path(*parts) if parts else Path() +def _has_parent_segment(raw: str, path: Path) -> bool: + """Return true when a user path contains a parent-directory segment. + + On POSIX, ``Path("E:\\foo\\..\\bar")`` treats backslashes as normal + characters, so check both the host parser and Windows-style parsing. + """ + if ".." in path.parts: + return True + if ".." in PureWindowsPath(raw).parts: + return True + return ".." in raw.replace("\\", "/").split("/") + + +def _is_absolute_user_path(path: Path) -> bool: + expanded = str(path) + if os.name == "nt": + return path.is_absolute() and PureWindowsPath(expanded).is_absolute() + return path.is_absolute() and PurePosixPath(expanded).is_absolute() + + def _assert_contained(resolved: Path, root: Path) -> None: """Raise ValueError if ``resolved`` realpaths outside ``root``.""" try: @@ -351,10 +371,10 @@ def resolve_under_root( raise ValueError("path may not contain null bytes") path = Path(raw).expanduser() - if ".." in path.parts: + if _has_parent_segment(raw, path): raise ValueError(f"path may not contain '..' segments: {raw!r}") - if path.is_absolute(): + if _is_absolute_user_path(path): _assert_contained(path, root) return path @@ -373,6 +393,36 @@ def resolve_output_dir(path_value: str | None = None) -> Path: def resolve_export_dir(path_value: str | None = None) -> Path: + """Resolve an export directory — contained under exports_root(). + + Used by scan/read endpoints. Use :func:`resolve_export_write_dir` + for the export write path where absolute paths are accepted. + """ + return resolve_under_root( + path_value, + root = exports_root(), + strip_prefixes = ("exports",), + ) + + +def resolve_export_write_dir(path_value: str | None = None) -> Path: + """Resolve an export save directory — accepts absolute paths. + + Unlike :func:`resolve_export_dir`, this function passes absolute + paths through as-is so users can target a different drive when + their Studio install lives on a constrained system volume + (see :gh-issue:`6082`). Used only by the export write path. + """ + if not path_value or not str(path_value).strip(): + return exports_root() + raw = str(path_value).strip() + if "\x00" in raw: + raise ValueError("path may not contain null bytes") + path = Path(raw).expanduser() + if _has_parent_segment(raw, path): + raise ValueError(f"path may not contain '..' segments: {raw!r}") + if _is_absolute_user_path(path): + return path return resolve_under_root( path_value, root = exports_root(), diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py index 98c48fe45c..9c18070fbb 100644 --- a/studio/backend/utils/studio_version.py +++ b/studio/backend/utils/studio_version.py @@ -15,6 +15,7 @@ _DEV_VERSION = "dev" _GIT_TIMEOUT_SECONDS = 1.0 _STUDIO_TAG_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$") _GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$") +_GIT_BRANCH_RE = re.compile(r"^[0-9A-Za-z._/-]+$") _MAX_VERSION_LENGTH = 64 @@ -71,6 +72,35 @@ def _exact_git_studio_tag(repo_root: Path) -> str | None: return tag if is_valid_studio_release_version(tag) else None +def _git_branch(repo_root: Path) -> str | None: + try: + result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + cwd = repo_root, + check = False, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = _GIT_TIMEOUT_SECONDS, + ) + except (OSError, subprocess.TimeoutExpired): + return None + + if result.returncode != 0: + return None + + branch = result.stdout.strip() + # "HEAD" means detached, e.g. a tag or commit checkout. + if ( + not branch + or branch == "HEAD" + or len(branch) > _MAX_VERSION_LENGTH + or _GIT_BRANCH_RE.fullmatch(branch) is None + ): + return None + return branch + + def get_studio_version(repo_root: Path | None = None) -> str: """Return the installed Studio release tag for display, or ``dev``. @@ -81,7 +111,10 @@ def get_studio_version(repo_root: Path | None = None) -> str: if _is_source_checkout(resolved_repo_root): git_tag = _exact_git_studio_tag(resolved_repo_root) - return git_tag if git_tag is not None else _DEV_VERSION + if git_tag is not None: + return git_tag + branch = _git_branch(resolved_repo_root) + return f"GitHub {branch}" if branch is not None else _DEV_VERSION stamped_version = _studio_release_build.STUDIO_RELEASE_VERSION if is_valid_studio_release_version(stamped_version): diff --git a/studio/frontend/index.html b/studio/frontend/index.html index 4f81ffd4ff..0fbb4eaeeb 100644 --- a/studio/frontend/index.html +++ b/studio/frontend/index.html @@ -1,16 +1,16 @@ - + - - - - - - Unsloth Studio - - -
- - - + + + + + + Unsloth Studio + + +
+ + + diff --git a/studio/frontend/public/hub/profile/logo/meta.svg b/studio/frontend/public/hub/profile/logo/meta.svg index 9fa656bd6b..fe3709aeea 100644 --- a/studio/frontend/public/hub/profile/logo/meta.svg +++ b/studio/frontend/public/hub/profile/logo/meta.svg @@ -1,19 +1,19 @@ - - -Logo of Meta Platforms -- Graphic created by Detmar Owen - - - - - - - - - - - - - - - + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/meta.svg b/studio/frontend/public/provider-logos/misc/meta.svg index 9fa656bd6b..fe3709aeea 100644 --- a/studio/frontend/public/provider-logos/misc/meta.svg +++ b/studio/frontend/public/provider-logos/misc/meta.svg @@ -1,19 +1,19 @@ - - -Logo of Meta Platforms -- Graphic created by Detmar Owen - - - - - - - - - - - - - - - + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 91d99e238e..f806a0f405 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -259,8 +259,16 @@ function TauriWrapper({ children }: { children: ReactNode }) { <> {children} - - +
+ + +
); } diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index c05b0e6451..f5179e68f7 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,7 +6,7 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; -import { useChatRuntimeStore } from "@/features/chat"; +import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat"; import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; import { useT, type TranslationKey } from "@/i18n"; @@ -15,6 +15,7 @@ import { createRootRoute, redirect, useMatches, + useNavigate, useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; @@ -78,6 +79,7 @@ function RootLayout() { const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname); const isChatRoute = pathname.startsWith("/chat"); const { pinned, setPinned, togglePinned } = useSidebarPin(); + const navigate = useNavigate(); useTrainingUnloadGuard(); @@ -107,11 +109,24 @@ function RootLayout() { if ((e.metaKey || e.ctrlKey) && e.key === ",") { e.preventDefault(); useSettingsDialogStore.getState().openDialog(); + return; + } + // Cmd/Ctrl+Shift+O opens a new chat. + if ((e.metaKey || e.ctrlKey) && e.shiftKey && e.code === "KeyO") { + e.preventDefault(); + clearNewChatDraft(); // fresh chat starts empty, no bleed from the last one + const chatRuntime = useChatRuntimeStore.getState(); + chatRuntime.setActiveThreadId(null); + chatRuntime.setActiveProjectId(null); + void navigate({ + to: "/chat", + search: { new: crypto.randomUUID() }, + }); } }; window.addEventListener("keydown", handler); return () => window.removeEventListener("keydown", handler); - }, []); + }, [navigate]); useEffect(() => { if (isChatRoute) return; diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 7baa69571c..16ab7932b1 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -45,6 +45,8 @@ import { Switch } from "@/components/ui/switch"; import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler"; import { cn } from "@/lib/utils"; import { + Archive01Icon, + ArchiveRestoreIcon, ChefHatIcon, CursorInfo02Icon, DashboardCircleIcon, @@ -60,6 +62,7 @@ import { Logout05Icon, MoreVerticalIcon, Search01Icon, + PlusSignIcon, PowerIcon, PencilEdit02Icon, LayoutAlignLeftIcon, @@ -82,13 +85,16 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { ChevronDown, ChevronsUpDown, MoreHorizontalIcon, Moon, Sun } from "lucide-react"; import { Link, useNavigate, useRouterState } from "@tanstack/react-router"; import { + archiveChatItem, ChatSearchDialog, + clearNewChatDraft, createChatProject, deleteChatProject, deleteChatItem, moveChatItemToProject, renameChatItem, renameChatProject, + unarchiveChatItem, useChatRuntimeStore, useChatProjects, useChatSearchStore, @@ -290,15 +296,16 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ - enabled: !isStudioRoute, - requireMessages: false, - }); + const { items: allChatItems, archivedItems: archivedChatItems } = + useChatSidebarItems({ + enabled: !isStudioRoute, + requireMessages: false, + }); const recentChatItems = useMemo( () => allChatItems.filter((item) => !item.projectId), [allChatItems], ); - const chatItems = allChatItems; + const [archivedOpen, setArchivedOpen] = useState(false); const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId); const activeThreadId = isChatRoute @@ -348,6 +355,7 @@ export function AppSidebar() { function openNewChat(projectId = activeProjectId) { if (chatDisabled) return; + clearNewChatDraft(); setActiveThreadId(null); useChatRuntimeStore.getState().setActiveProjectId(projectId); navigate({ to: "/chat", search: chatSearchForProject(projectId) }); @@ -373,6 +381,33 @@ export function AppSidebar() { }); } + async function handleArchiveThread(item: SidebarItem) { + try { + await archiveChatItem(item, activeThreadId, (view) => { + navigate({ + to: "/chat", + search: item.projectId + ? { project: item.projectId } + : { new: view.newThreadNonce }, + }); + }); + } catch (err) { + toast.error("Failed to archive chat", { + description: err instanceof Error ? err.message : undefined, + }); + } + } + + async function handleUnarchiveThread(item: SidebarItem) { + try { + await unarchiveChatItem(item); + } catch (err) { + toast.error("Failed to unarchive chat", { + description: err instanceof Error ? err.message : undefined, + }); + } + } + type RenameTarget = | { kind: "chat"; item: SidebarItem; current: string } | { kind: "project"; project: ProjectRecord; current: string } @@ -558,7 +593,8 @@ export function AppSidebar() { : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; const buttonClass = cn( "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium", - variant === "project" ? "pl-[39px]" : "pl-3", + // pl-3.5 starts the title at the same x as the Recents label text. + variant === "project" ? "pl-[39px]" : "pl-3.5", variant === "project" ? "group-hover/project-chat-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8" : "group-hover/recent-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8", @@ -608,7 +644,7 @@ export function AppSidebar() { side="bottom" align="start" sideOffset={0} - className="unsloth-plus-menu menu-flat-destructive w-52" + className="unsloth-plus-menu menu-flat-destructive w-56" > openRenameChat(item)}> @@ -690,6 +726,10 @@ export function AppSidebar() { + void handleArchiveThread(item)}> + + Archive + setConfirmingDelete({ kind: "chat", item })} @@ -837,7 +877,28 @@ export function AppSidebar() { navigate({ to: "/projects" }); closeMobileIfOpen(); }} - /> + className="group/projects-item relative" + > + + - + {recentChatItems.map((item) => renderChatSidebarItem(item, "recent"), @@ -938,6 +999,85 @@ export function AppSidebar() { )} + {/* Archived chats — hidden on Studio + when nothing is archived */} + {!isStudioRoute && archivedChatItems.length > 0 && ( + + + + + Archived + + + + + + + {archivedChatItems.map((item) => ( + + { + navigate({ + to: "/chat", + search: + item.type === "single" + ? { thread: item.id } + : { compare: item.id }, + }); + closeMobileIfOpen(); + }} + > + {item.title} + + + + + + + openRenameChat(item)}> + + Rename + + void handleUnarchiveThread(item)}> + + Unarchive + + setConfirmingDelete({ kind: "chat", item })} + > + + Delete + + + + + ))} + + + + + + )} + {isStudioRoute && runItems.length > 0 && !chatOnly && ( @@ -948,7 +1088,7 @@ export function AppSidebar() { - + {runItems.map((run) => { // Explicit selection wins. Otherwise highlight the active @@ -1065,7 +1205,7 @@ export function AppSidebar() { className="!size-[32px]" /> -
+
{displayTitle} Unsloth
@@ -1076,7 +1216,7 @@ export function AppSidebar() { side="top" align="center" sideOffset={8} - className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-1.5 py-2.5 font-heading rounded-[20px] border-0" + className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-2.5 py-2.5 font-heading rounded-[20px] border-0" > state.artifactsEnabled || state.collapseHtmlArtifacts, diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 447fd763e3..8d40f587ad 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -19,6 +19,11 @@ const formatNumber = (n: number): string => { return n.toLocaleString(); }; +const formatRate = (r: number | undefined): string => { + if (r === undefined || !Number.isFinite(r)) return "—"; + return `${Math.round(r).toLocaleString()} tok/s`; +}; + /** * Shows streaming stats as a badge with hover tooltip. * When server timings are available (GGUF), shows prompt eval, generation, @@ -51,6 +56,10 @@ export const MessageTiming: FC<{ st?.cache_n ?? custom?.contextUsage?.cachedTokens ?? 0; // Anthropic-only cache-write count. const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0; + // DiffusionGemma reports separately-labelled throughput (no prefill, so no "prompt + // speed"), matching the CLI: in-step parallel, effective (canvas*blocks/wall), and + // output (answer tokens/wall). + const isDiffusion = (st as { diffusion?: boolean } | undefined)?.diffusion === true; // Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op turns, // blowing the rate up to Infinity. Require >=1 token, a non-zero decode @@ -93,6 +102,99 @@ export const MessageTiming: FC<{ >
{st ? ( + isDiffusion ? ( + <> + {/* DiffusionGemma: honest throughput (no autoregressive prompt speed) */} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} + {st?.diffusion_parallel_tok_s != null && ( +
+ Speed (in-step) + + {formatRate(st.diffusion_parallel_tok_s)} + +
+ )} + {st?.diffusion_effective_tok_s != null && ( +
+ Effective + + {formatRate(st.diffusion_effective_tok_s)} + +
+ )} + {st?.diffusion_output_tok_s != null && ( +
+ Output + + {formatRate(st.diffusion_output_tok_s)} + +
+ )} + {st?.diffusion_steps != null && ( +
+ Denoising + + {formatNumber(st.diffusion_steps)} steps + {st?.diffusion_blocks != null + ? `, ${formatNumber(st.diffusion_blocks)} block${st.diffusion_blocks === 1 ? "" : "s"}` + : ""} + +
+ )} + {st?.diffusion_canvas != null && ( +
+ Canvas + + {formatNumber(st.diffusion_canvas)} tokens + +
+ )} + {(st?.diffusion_wall_ms ?? st?.predicted_ms) != null && ( +
+ Generation + + {formatTimingMs(st.diffusion_wall_ms ?? st.predicted_ms)} + +
+ )} + {timing.tokenCount !== undefined && ( +
+ Answer tokens + + {formatNumber(timing.tokenCount)} + +
+ )} + {(st?.diffusion_prompt_n ?? st?.prompt_n) != null && ( +
+ Prompt + + {formatNumber(st.diffusion_prompt_n ?? st.prompt_n)} tokens + +
+ )} +
+
+ Total + + {formatTimingMs(timing.totalStreamTime)} + +
+
+ Chunks + + {timing.totalChunks} + +
+ + ) : ( <> {/* Server-side metrics (GGUF) */} {st?.prompt_ms != null && ( @@ -135,6 +237,30 @@ export const MessageTiming: FC<{
)} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} + {st?.diffusion_steps != null && ( +
+ Denoising steps + + {formatNumber(st.diffusion_steps)} + +
+ )} + {st?.diffusion_blocks != null && ( +
+ Blocks + + {formatNumber(st.diffusion_blocks)} + +
+ )} {cacheHits > 0 && (
Cache hits @@ -165,6 +291,7 @@ export const MessageTiming: FC<{
+ ) ) : ( <> {/* Client-side metrics (safetensors + external provider fallback) */} diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index b8d1bc0f00..fa6980824f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -189,7 +189,7 @@ function ModelSelectorTrigger({ @@ -247,7 +247,7 @@ function ModelSelectorContent({ alignOffset={10} data-tour={dataTour} className={cn( - "unsloth-model-selector-menu menu-soft-surface ring-0 w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 p-3", + "unsloth-model-selector-menu menu-soft-surface ring-0 w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 px-3 pt-3 pb-2", className, )} > @@ -307,7 +307,7 @@ function ModelSelectorContent({ )} {onPickLocalModel ? ( -
+
+ + + {failed ? ( + + Could not send your decision. Try again. + + ) : null} +
+ ); +} + +export function withToolConfirmation( + Component: ToolCallMessagePartComponent, +): ToolCallMessagePartComponent { + const WithToolConfirmation: ToolCallMessagePartComponent = (props) => ( + <> + + + + ); + return WithToolConfirmation; +} diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index eeb4c2059f..8930b6386f 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -325,6 +325,9 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ result, status, }) => { + // Allow/Deny confirmation controls are rendered uniformly for every tool + // card (built-in and fallback) by the `withToolConfirmation` wrapper in + // thread.tsx, so this renderer stays purely presentational. const isCancelled = status?.type === "incomplete" && status.reason === "cancelled"; diff --git a/studio/frontend/src/components/assistant-ui/tool-group.tsx b/studio/frontend/src/components/assistant-ui/tool-group.tsx index 569e673da3..3ce65627e2 100644 --- a/studio/frontend/src/components/assistant-ui/tool-group.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-group.tsx @@ -9,6 +9,7 @@ import { type PropsWithChildren, } from "react"; import { useAuiState } from "@assistant-ui/react"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { ChevronDownIcon } from "lucide-react"; import { Wrench01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -216,6 +217,31 @@ const ToolGroupImpl: FC< (part) => part.type === "tool-call" && part.toolName === "render_html", ), ); + // A blocking allow/deny prompt must never be hidden inside a collapsed + // group, so force the group open while any of its calls awaits confirmation. + const toolConfirmations = useChatRuntimeStore((s) => s.toolConfirmations); + const hasPendingConfirmation = useAuiState(({ message }) => + message.parts + .slice(startIndex, endIndex + 1) + .some( + (part) => + part.type === "tool-call" && + Object.prototype.hasOwnProperty.call( + toolConfirmations, + part.toolCallId, + ), + ), + ); + const messageRunning = useAuiState( + ({ message }) => message.status?.type === "running", + ); + // Keep the group open once a confirmation forced it open, so answering an + // allow/deny doesn't snap it shut between sequential tool calls. It reverts + // to the default collapsed state once the turn finishes. + const forcedOpenRef = useRef(false); + if (hasPendingConfirmation) forcedOpenRef.current = true; + const forceOpen = + hasPendingConfirmation || (forcedOpenRef.current && messageRunning); // Render single tool calls and artifacts directly so cards never hide in a // collapsed group. @@ -224,7 +250,7 @@ const ToolGroupImpl: FC< } return ( - + {children} diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index c8a4dd1e40..88826e57e5 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -3,27 +3,100 @@ import { Button } from "@/components/ui/button"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; +import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; import { AnimatePresence, motion } from "motion/react"; -import { type ReactElement, useEffect, useRef } from "react"; +import { type ReactElement, useEffect, useRef, useState } from "react"; const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1]; +// Backend progress is coarse (5% steps, ~0.9 max) and the extract tail emits no +// signal. Creep toward this cap so the bar keeps moving rather than freezing. +const RUNNING_CAP = 0.95; + +// Smoothed 0..1 bar progress: eases toward real `progress`, trickles toward a +// ceiling when idle, animates to 100% when `done`. Resets to 0 on each start. +function useSmoothedProgress( + active: boolean, + progress: number | null, + done: boolean, +): number { + const [display, setDisplay] = useState(0); + const displayRef = useRef(0); + const progressRef = useRef(progress); + const doneRef = useRef(done); + progressRef.current = progress; + doneRef.current = done; + + useEffect(() => { + if (!active) { + displayRef.current = 0; + setDisplay(0); + return; + } + let raf = 0; + let last = performance.now(); + const tick = (now: number) => { + // rAF timestamps can predate the performance.now() captured above, so + // clamp dt at 0 to keep the first frame from stepping backwards. + const dt = Math.max(0, Math.min((now - last) / 1000, 0.1)); + last = now; + const current = displayRef.current; + const real = progressRef.current ?? 0; + let target: number; + let speed: number; // approach rate (fraction of remaining gap per second) + if (doneRef.current) { + target = 1; + speed = 5; + } else if (real > current) { + target = real; // catch up to a freshly observed milestone + speed = 4; + } else { + target = RUNNING_CAP; // no signal: creep toward the cap, never frozen + speed = 0.3; + } + const cap = doneRef.current ? 1 : RUNNING_CAP; + const next = Math.min( + current + (target - current) * Math.min(speed * dt, 1), + cap, + ); + displayRef.current = next; + setDisplay(next); + if (doneRef.current && next > 0.999) { + return; + } + raf = requestAnimationFrame(tick); + }; + raf = requestAnimationFrame(tick); + return () => cancelAnimationFrame(raf); + }, [active]); + + return display; +} interface LlamaUpdateBannerProps { enabled?: boolean; + // false: fill the parent instead of self-anchoring, so banners can stack in a + // shared container. true (default) keeps standalone desktop mounts working. + positioned?: boolean; } /** * Non-invasive "Update llama.cpp" affordance. Appears bottom-right ~1s after a - * newer prebuilt is detected and stays up until dismissed (click outside / X) - * or updated. Clicking Update swaps the prebuilt in place via POST /api/llama/update. + * newer prebuilt is detected and stays up until the user explicitly acts on it + * (X, Update, or Remind me later). Clicking Update swaps the prebuilt in place + * via POST /api/llama/update. Can be turned off entirely in Settings -> + * General -> Notifications (on by default). */ export function LlamaUpdateBanner({ enabled = true, + positioned = true, }: LlamaUpdateBannerProps): ReactElement | null { - const { status, visible, applying, apply, dismiss } = useLlamaUpdateCheck({ - enabled, - }); + const showBannerPref = useShowLlamaUpdateBanner(); + const { status, visible, applying, apply, dismiss, snooze } = + useLlamaUpdateCheck({ + enabled: enabled && showBannerPref, + }); async function handleUpdate() { const result = await apply(); @@ -40,43 +113,36 @@ export function LlamaUpdateBanner({ const show = visible && status != null && (status.update_available || applying); - const bannerRef = useRef(null); - - // Dismiss when the user clicks anything outside the banner. Kept off while an - // update is applying so the progress stays visible. - useEffect(() => { - if (!show || applying) return; - function onPointerDown(event: PointerEvent) { - if ( - bannerRef.current && - !bannerRef.current.contains(event.target as Node) - ) { - dismiss(); - } - } - document.addEventListener("pointerdown", onPointerDown, true); - return () => - document.removeEventListener("pointerdown", onPointerDown, true); - }, [show, applying, dismiss]); + const updateProgress = status?.job.progress ?? null; + const jobSucceeded = status?.job.state === "success"; + // Drives the bar so it animates continuously; aria reports the real value. + const displayProgress = useSmoothedProgress( + applying, + updateProgress, + jobSucceeded, + ); return ( {show ? ( -
+
{applying ? null : ( -
+
+
+ ) : ( +
+ + +
+ )}
) : null} diff --git a/studio/frontend/src/components/tauri/startup-screen.tsx b/studio/frontend/src/components/tauri/startup-screen.tsx index 1eb3f81d10..8e373c82a0 100644 --- a/studio/frontend/src/components/tauri/startup-screen.tsx +++ b/studio/frontend/src/components/tauri/startup-screen.tsx @@ -228,7 +228,7 @@ function RepairingContent({
-

Updating existing Studio install...

+

Updating existing Unsloth install...

{latest && (

{latest}

)} diff --git a/studio/frontend/src/components/ui/accordion.tsx b/studio/frontend/src/components/ui/accordion.tsx index 7754c78a11..35de233858 100644 --- a/studio/frontend/src/components/ui/accordion.tsx +++ b/studio/frontend/src/components/ui/accordion.tsx @@ -1,98 +1,98 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"use client"; - -import { Accordion as AccordionPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; -import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; - -function Accordion({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AccordionItem({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AccordionTrigger({ - className, - children, - ...props -}: React.ComponentProps) { - return ( - - - {children} - - - - - ); -} - -function AccordionContent({ - className, - children, - ...props -}: React.ComponentProps) { - return ( - -
- {children} -
- - ); -} - -export { Accordion, AccordionItem, AccordionTrigger, AccordionContent }; +"use client"; + +import { Accordion as AccordionPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; +import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +function Accordion({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AccordionItem({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AccordionTrigger({ + className, + children, + ...props +}: React.ComponentProps) { + return ( + + + {children} + + + + + ); +} + +function AccordionContent({ + className, + children, + ...props +}: React.ComponentProps) { + return ( + +
+ {children} +
+ + ); +} + +export { Accordion, AccordionItem, AccordionTrigger, AccordionContent }; diff --git a/studio/frontend/src/components/ui/alert-dialog.tsx b/studio/frontend/src/components/ui/alert-dialog.tsx index f5c1dbacca..97f4be7f44 100644 --- a/studio/frontend/src/components/ui/alert-dialog.tsx +++ b/studio/frontend/src/components/ui/alert-dialog.tsx @@ -1,50 +1,50 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { AlertDialog as AlertDialogPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { Button } from "@/components/ui/button"; -import { cn } from "@/lib/utils"; - -function AlertDialog({ - ...props -}: React.ComponentProps) { - return ; -} - -function AlertDialogTrigger({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogPortal({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogOverlay({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - +import { AlertDialog as AlertDialogPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; + +function AlertDialog({ + ...props +}: React.ComponentProps) { + return ; +} + +function AlertDialogTrigger({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogPortal({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogOverlay({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + function AlertDialogContent({ className, size = "default", @@ -60,143 +60,143 @@ function AlertDialogContent({ - - ); -} - -function AlertDialogHeader({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogFooter({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogMedia({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogTitle({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogDescription({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogAction({ - className, - variant = "default", - size = "default", - ...props -}: React.ComponentProps & - Pick, "variant" | "size">) { - return ( - - ); -} - -function AlertDialogCancel({ - className, - variant = "outline", - size = "default", - ...props -}: React.ComponentProps & - Pick, "variant" | "size">) { - return ( - - ); -} - -export { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogMedia, - AlertDialogOverlay, - AlertDialogPortal, - AlertDialogTitle, - AlertDialogTrigger, -}; + className={cn( + "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 bg-background ring-foreground/5 gap-6 rounded-4xl p-6 ring-1 duration-100 data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-md group/alert-dialog-content fixed top-1/2 left-1/2 z-50 grid w-full -translate-x-1/2 -translate-y-1/2 outline-none", + className, + )} + {...props} + /> + + ); +} + +function AlertDialogHeader({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogFooter({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogMedia({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogTitle({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogDescription({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogAction({ + className, + variant = "default", + size = "default", + ...props +}: React.ComponentProps & + Pick, "variant" | "size">) { + return ( + + ); +} + +function AlertDialogCancel({ + className, + variant = "outline", + size = "default", + ...props +}: React.ComponentProps & + Pick, "variant" | "size">) { + return ( + + ); +} + +export { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogMedia, + AlertDialogOverlay, + AlertDialogPortal, + AlertDialogTitle, + AlertDialogTrigger, +}; diff --git a/studio/frontend/src/components/ui/alert.tsx b/studio/frontend/src/components/ui/alert.tsx index a4a5f4c4b7..094b607d9a 100644 --- a/studio/frontend/src/components/ui/alert.tsx +++ b/studio/frontend/src/components/ui/alert.tsx @@ -1,79 +1,79 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { type VariantProps, cva } from "class-variance-authority"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -const alertVariants = cva( - "grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert", - { - variants: { - variant: { - default: "bg-card text-card-foreground", - destructive: - "text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current", - }, - }, - defaultVariants: { - variant: "default", - }, - }, -); - -function Alert({ - className, - variant, - ...props -}: React.ComponentProps<"div"> & VariantProps) { - return ( -
- ); -} - -function AlertTitle({ className, ...props }: React.ComponentProps<"div">) { - return ( -
svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3", - className, - )} - {...props} - /> - ); -} - -function AlertDescription({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertAction({ className, ...props }: React.ComponentProps<"div">) { - return ( -
- ); -} - -export { Alert, AlertTitle, AlertDescription, AlertAction }; +import { type VariantProps, cva } from "class-variance-authority"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +const alertVariants = cva( + "grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert", + { + variants: { + variant: { + default: "bg-card text-card-foreground", + destructive: + "text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current", + }, + }, + defaultVariants: { + variant: "default", + }, + }, +); + +function Alert({ + className, + variant, + ...props +}: React.ComponentProps<"div"> & VariantProps) { + return ( +
+ ); +} + +function AlertTitle({ className, ...props }: React.ComponentProps<"div">) { + return ( +
svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3", + className, + )} + {...props} + /> + ); +} + +function AlertDescription({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertAction({ className, ...props }: React.ComponentProps<"div">) { + return ( +
+ ); +} + +export { Alert, AlertTitle, AlertDescription, AlertAction }; diff --git a/studio/frontend/src/components/ui/animated-shiny-text.tsx b/studio/frontend/src/components/ui/animated-shiny-text.tsx index 4c650f1003..8d366ca3d6 100644 --- a/studio/frontend/src/components/ui/animated-shiny-text.tsx +++ b/studio/frontend/src/components/ui/animated-shiny-text.tsx @@ -1,41 +1,41 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react" - -import { cn } from "@/lib/utils" - -export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> { - shimmerWidth?: number -} - -export const AnimatedShinyText: FC = ({ - children, - className, - shimmerWidth = 100, - ...props -}) => { - return ( - - {children} - - ) -} +import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react" + +import { cn } from "@/lib/utils" + +export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> { + shimmerWidth?: number +} + +export const AnimatedShinyText: FC = ({ + children, + className, + shimmerWidth = 100, + ...props +}) => { + return ( + + {children} + + ) +} diff --git a/studio/frontend/src/components/ui/aspect-ratio.tsx b/studio/frontend/src/components/ui/aspect-ratio.tsx index cb605f01eb..2471f4333d 100644 --- a/studio/frontend/src/components/ui/aspect-ratio.tsx +++ b/studio/frontend/src/components/ui/aspect-ratio.tsx @@ -1,12 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { AspectRatio as AspectRatioPrimitive } from "radix-ui"; - -function AspectRatio({ - ...props -}: React.ComponentProps) { - return ; -} - -export { AspectRatio }; +import { AspectRatio as AspectRatioPrimitive } from "radix-ui"; + +function AspectRatio({ + ...props +}: React.ComponentProps) { + return ; +} + +export { AspectRatio }; diff --git a/studio/frontend/src/components/ui/avatar.tsx b/studio/frontend/src/components/ui/avatar.tsx index 2250bb849a..31262b32f7 100644 --- a/studio/frontend/src/components/ui/avatar.tsx +++ b/studio/frontend/src/components/ui/avatar.tsx @@ -1,113 +1,113 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Avatar as AvatarPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -function Avatar({ - className, - size = "default", - ...props -}: React.ComponentProps & { - size?: "default" | "sm" | "lg"; -}) { - return ( - - ); -} - -function AvatarImage({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AvatarFallback({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) { - return ( - svg]:hidden", - "group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2", - "group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2", - className, - )} - {...props} - /> - ); -} - -function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AvatarGroupCount({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2", - className, - )} - {...props} - /> - ); -} - -export { - Avatar, - AvatarImage, - AvatarFallback, - AvatarGroup, - AvatarGroupCount, - AvatarBadge, -}; +import { Avatar as AvatarPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +function Avatar({ + className, + size = "default", + ...props +}: React.ComponentProps & { + size?: "default" | "sm" | "lg"; +}) { + return ( + + ); +} + +function AvatarImage({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AvatarFallback({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) { + return ( + svg]:hidden", + "group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2", + "group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2", + className, + )} + {...props} + /> + ); +} + +function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AvatarGroupCount({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2", + className, + )} + {...props} + /> + ); +} + +export { + Avatar, + AvatarImage, + AvatarFallback, + AvatarGroup, + AvatarGroupCount, + AvatarBadge, +}; diff --git a/studio/frontend/src/components/ui/badge.tsx b/studio/frontend/src/components/ui/badge.tsx index 3951ae9de0..0f2f334986 100644 --- a/studio/frontend/src/components/ui/badge.tsx +++ b/studio/frontend/src/components/ui/badge.tsx @@ -1,54 +1,54 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -/* eslint-disable react-refresh/only-export-components */ - -import { type VariantProps, cva } from "class-variance-authority"; -import { Slot } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -export const badgeVariants = cva( - "h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge", - { - variants: { - variant: { - default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", - secondary: - "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", - destructive: - "bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20", - outline: - "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground bg-input/30", - ghost: - "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", - link: "text-primary underline-offset-4 hover:underline", - }, - }, - defaultVariants: { - variant: "default", - }, - }, -); - -export function Badge({ - className, - variant = "default", - asChild = false, - ...props -}: React.ComponentProps<"span"> & - VariantProps & { - asChild?: boolean; - }): React.ReactElement { - const Comp = asChild ? Slot.Root : "span"; - - return ( - - ); -} +/* eslint-disable react-refresh/only-export-components */ + +import { type VariantProps, cva } from "class-variance-authority"; +import { Slot } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +export const badgeVariants = cva( + "h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge", + { + variants: { + variant: { + default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", + secondary: + "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", + destructive: + "bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20", + outline: + "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground bg-input/30", + ghost: + "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", + link: "text-primary underline-offset-4 hover:underline", + }, + }, + defaultVariants: { + variant: "default", + }, + }, +); + +export function Badge({ + className, + variant = "default", + asChild = false, + ...props +}: React.ComponentProps<"span"> & + VariantProps & { + asChild?: boolean; + }): React.ReactElement { + const Comp = asChild ? Slot.Root : "span"; + + return ( + + ); +} diff --git a/studio/frontend/src/components/ui/breadcrumb.tsx b/studio/frontend/src/components/ui/breadcrumb.tsx index dc026994ce..a2dad8783f 100644 --- a/studio/frontend/src/components/ui/breadcrumb.tsx +++ b/studio/frontend/src/components/ui/breadcrumb.tsx @@ -1,126 +1,126 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Slot } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; -import { - ArrowRight01Icon, - MoreHorizontalCircle01Icon, -} from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; - -function Breadcrumb({ className, ...props }: React.ComponentProps<"nav">) { - return ( -