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

Merge remote-tracking branch 'origin/main' into docker-blackwell-build

Browse source
This commit is contained in:
Daniel Han 2026-07-08 06:46:38 +00:00
commit d34cb71189
115 changed files with 23615 additions and 1054 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

2
.gitignore vendored
View file

@ -11,6 +11,8 @@ outputs/
exports/
/datasets/
studio/backend/assets/datasets/
# Generated async worker / reviewer transcripts (never part of the product).
studio/backend/async_task_outputs/
unsloth_training_checkpoints/
*.gguf
*.safetensors

10
install.ps1
View file

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

182
install.sh
View file

@ -1442,8 +1442,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then
SKIP_TORCH=true
fi
# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for
# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file
# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover.
_MLX_LM_EXCLUDE_ARG=""
# Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file).
if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3"
_OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt"
if [ -f "$_OVERRIDES_FILE" ]; then
# uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace
@ -1477,6 +1483,81 @@ elif [ "$OS" = "macos" ]; then
fi
tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none"
# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in
# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded
# to 10s. Defined here so the reroute below can use it before _run_bounded exists.
_WSL_AMD_GPU_NAME_CACHE=""
_wsl_amd_gpu_name() {
if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then
[ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1
printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0
fi
command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; }
_wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name"
if command -v timeout >/dev/null 2>&1; then
_wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
else
_wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
fi
if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi
_WSL_AMD_GPU_NAME_CACHE="-"; return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04)
# with a 24.04 distro present, re-run the install there and stop; else fall through
# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before
@ -1487,7 +1568,15 @@ _maybe_reroute_strixhalo_to_2404() {
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
[ -e /dev/dxg ] || return 0
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on
# this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
if _has_usable_nvidia_gpu; then return 0; fi
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
# Already ROCm-on-WSL? leave a working GPU alone, whatever the version.
if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then
return 0
@ -1963,61 +2052,6 @@ _has_amd_rocm_gpu() {
return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@ -2341,19 +2375,19 @@ _persist_rocm_wsl_dropin() {
fi
}
# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it.
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
# Leave any already-usable GPU completely alone (NVIDIA, or working ROCm).
if _has_usable_nvidia_gpu; then return 0; fi
# "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the
# generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and
# would skip this bootstrap while the real GPU is still unusable. awk consumes
# all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail.
# Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000,
# the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so
# rocminfo isn't SIGPIPE'd like `grep -q` under pipefail.
_ensure_rocm_probe_env
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then
# rocminfo may work only via the transient env _ensure_rocm_probe_env
# just set, which dies with the installer. Persist the drop-in so login
# shells (Studio, llama.cpp) inherit it -- else a reinstall over an
@ -2363,9 +2397,12 @@ _maybe_bootstrap_rocm_wsl() {
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
[ -e /dev/dxg ] || return 0
# Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match
# the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S").
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
command -v bash >/dev/null 2>&1 || return 0
# Fast path: already configured (librocdxg present) but launched from a
@ -2383,7 +2420,8 @@ _maybe_bootstrap_rocm_wsl() {
fi
echo ""
substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN"
_rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU"
substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN"
substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU."
substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)"
@ -2688,7 +2726,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2699,9 +2737,11 @@ if [ "$_MIGRATED" = true ]; then
run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
fi
else
# Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2905,7 +2945,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2923,7 +2963,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
--upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git ${_ZOO_REF}..."
@ -2932,7 +2972,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
"$_ZOO_GIT_SPEC"
else
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth -- "$PACKAGE_NAME"
--upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-}
fi
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
@ -2955,7 +2995,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git ${_ZOO_REF}..."

6
pyproject.toml
View file

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

57
scripts/install_rocm_wsl_strixhalo.sh
View file

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

72
scripts/scan_packages_baseline.json
View file

@ -1489,6 +1489,78 @@
"severity": "HIGH",
"evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2",
"evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381"
},
{
"package": "fastapi",
"file": "fastapi/routing.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45",
"evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5"
},
{
"package": "fastmcp-slim",
"file": "fastmcp/cli/apps_dev.py",
"check": "Enumerates filesystem AND makes network calls",
"severity": "CRITICAL",
"evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:",
"evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4"
},
{
"package": "huggingface-hub",
"file": "huggingface_hub/_sandbox.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38",
"evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8"
},
{
"package": "huggingface-hub",
"file": "huggingface_hub/_sandbox.py",
"check": "Writes to /tmp and executes (staged dropper)",
"severity": "CRITICAL",
"evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"",
"evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22"
},
{
"package": "huggingface-hub",
"file": "huggingface_hub/hf_api.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
"evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
},
{
"package": "huggingface-hub",
"file": "huggingface_hub/utils/_http.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
"evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958",
"evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e"
},
{
"package": "cffi",
"file": "cffi/_cffi_gen_src.py",
"check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
"severity": "HIGH",
"evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)",
"evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9"
},
{
"package": "multiprocess",
"file": "multiprocess/forkserver.py",
"check": "Reverse shell / bind shell pattern",
"severity": "CRITICAL",
"evidence": "L6: import socket sha256:6c707119169286c9a798e2c8d13a48614e481d8a503950916fd4ffb4c94d3182",
"evidence_hash": "50fec0f0522a8e4e636bf348b752002d7935d8455af31fb78c6f11e2eba19f6d"
},
{
"package": "multiprocess",
"file": "multiprocess/tests/__init__.py",
"check": "Reverse shell / bind shell pattern",
"severity": "CRITICAL",
"evidence": "L3569: os.dup2(conn.fileno(), i) | L3601: \"test needs os.dup2()\") | L3619: os.dup2(fd, newfd) | L20: import socket sha256:c824dc0f409f242420c3fbb324790c53cb3078d2c8b07ee8f2a05694b01c2946",
"evidence_hash": "3878a2b430c175dbc5877a95195bfe52f9588ff73fb74e2261ed5e33087915ad"
}
]
}

20
scripts/verify_import_hoist.py
View file

@ -564,6 +564,12 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
for n, tids in b["module_import_targets"].items():
if tids & after_used:
continue # resolved -> fine
# `from __future__ import ...` is a compiler directive, not a runtime
# binding: the name (`annotations`, ...) is never loaded, so it can never
# "resolve" to a use. Skip it so a legitimately-added future import
# (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged.
if all(t.startswith("from:__future__:") for t in tids):
continue
newly_added = bool(tids - before_module_targets)
was_used_before = bool(tids & before_used)
if newly_added or was_used_before:
@ -588,9 +594,23 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
# package object and only *add* submodule attributes (e.g. adding
# `import urllib.error` next to `import urllib.request`). Nothing the name
# resolved to before is lost, so no reference is re-pointed -- skip it.
#
# A deliberate *relocation* is also benign and must not block: when a name
# keeps its spelling but its import source is moved A -> B in THIS diff (the
# old `from A import x` is removed at module level and a new `from B import x`
# is added), the swap is intentional, not a silent re-point to a pre-existing
# different object. This mirrors the relocation tolerance already applied to
# TARGET-MISSING. The dangerous case -- the name now resolving to a target
# that already existed before (shadow/clash) -- is NOT exempted.
removed_module_targets = before_module_targets - after_module_targets
for key, tafter in b["target_by_use"].items():
tbefore = a["target_by_use"].get(key)
if tbefore and tbefore != tafter and (tbefore - tafter):
lost = tbefore - tafter
gained = tafter - tbefore
relocated = lost <= removed_module_targets and gained <= added_module_targets
if relocated:
continue
findings.append(
(
"BLOCKER",

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

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

43
studio/backend/core/inference/__init__.py
View file

@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation.
The default get_inference_backend() returns an InferenceOrchestrator that
delegates to a subprocess. The original InferenceBackend runs inside the
subprocess and can be imported directly from .inference when needed.
Public names are resolved lazily (PEP 562): importing this package -- or a
dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull
the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML
backend and its Studio dependencies). Those load only when a public name is
actually accessed, so standalone helpers stay unit-testable without the full
inference stack.
"""
from .orchestrator import InferenceOrchestrator, get_inference_backend
from .llama_cpp import LlamaCppBackend
# Expose InferenceOrchestrator as InferenceBackend for backward compat.
InferenceBackend = InferenceOrchestrator
from typing import TYPE_CHECKING
__all__ = [
"InferenceBackend",
@ -21,3 +24,33 @@ __all__ = [
"get_inference_backend",
"LlamaCppBackend",
]
# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator.
_LAZY_ATTRS = {
"InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"),
"InferenceBackend": ("orchestrator", "InferenceOrchestrator"),
"get_inference_backend": ("orchestrator", "get_inference_backend"),
"LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"),
}
def __getattr__(name):
try:
submodule, attr = _LAZY_ATTRS[name]
except KeyError:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
from importlib import import_module
value = getattr(import_module(f"{__name__}.{submodule}"), attr)
globals()[name] = value # cache so later access skips __getattr__
return value
def __dir__():
return sorted(set(globals()) | set(__all__))
if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names
from .llama_cpp import LlamaCppBackend
from .orchestrator import InferenceOrchestrator, get_inference_backend
InferenceBackend = InferenceOrchestrator

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

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

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

@ -0,0 +1,109 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Resolve a chat model's assistant-turn-end stop tokens.
Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships
config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its
small chat variants ship no generation_config), so generation runs past the turn
and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns.
Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it
actually uses to end a turn), not raw vocab membership: a base/coder model can
carry ChatML control tokens in a shared vocab without using them, and a loader
may have synced ``eos_token`` to the document terminator. Dependency-light (no
torch / unsloth) so it is unit-testable without the full inference stack.
"""
from typing import Optional
# Canonical assistant-turn-end markers per chat family.
_CHAT_TURN_END_TOKENS = (
"<|im_end|>", # ChatML: Qwen, Yi
"<|eot_id|>", # Llama 3.x
"<|eom_id|>", # Llama 3.x tool turns
"<end_of_turn>", # Gemma
"<turn|>", # Gemma-4
"<|end|>", # Phi
"<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's)
)
# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has
# its own streamer, so its eos is left untouched.
_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>")
def _eos_id_set(eos_token_id) -> set:
if isinstance(eos_token_id, (list, tuple)):
return {int(t) for t in eos_token_id if t is not None}
if eos_token_id is not None:
return {int(eos_token_id)}
return set()
def _collect_template_text(chat_template) -> str:
"""Flatten a tokenizer ``chat_template`` into one scannable string.
Usually the template is a single jinja string, but multi-variant models
(e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a
``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list
of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case
would skip turn-end detection for those valid models, so gather every string
leaf (variant names are harmless: they never contain the markers).
"""
if isinstance(chat_template, str):
return chat_template
if isinstance(chat_template, dict):
values = chat_template.values()
elif isinstance(chat_template, (list, tuple)):
values = chat_template
else:
return ""
parts = [_collect_template_text(v) for v in values]
return "\n".join(p for p in parts if p)
def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list:
"""eos of ``id_tokenizer`` plus any canonical turn-end marker the
``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` --
the tokenizer generation actually uses.
Pass the same tokenizer for both at load time. After a mapped ``get_chat_template``
pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective
template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped
template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab
folds the turn-end token onto the doc-eos id, and generate_stream re-reads the
original tokenizer, so resolving ids on the mapped tokenizer would store the wrong
(doc-eos) id and let generation run past the real turn marker."""
ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None))
template = _collect_template_text(getattr(template_tokenizer, "chat_template", None))
if not template or any(h in template for h in _HARMONY_MARKERS):
return sorted(ids)
unk = getattr(id_tokenizer, "unk_token_id", None)
for marker in _CHAT_TURN_END_TOKENS:
if marker in template:
try:
tid = id_tokenizer.convert_tokens_to_ids(marker)
except Exception:
tid = None
if tid is not None and tid != unk and int(tid) >= 0:
ids.add(int(tid))
return sorted(ids)
def resolve_chat_turn_end_eos_ids(tokenizer) -> list:
"""tokenizer.eos plus any canonical turn-end marker the model's chat_template
actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended
to be resolved once at load. Returns eos unchanged for harmony templates."""
return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer)
def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]:
"""Merged eos_token_id list, or None if ``current_eos`` already covers every
resolved turn-end id. Used to repair a model's generation_config at load so
every ``.generate()`` path (vision, tool loops) stops at the turn boundary."""
if not turn_end_ids:
return None
current_set = _eos_id_set(current_eos)
if set(turn_end_ids) <= current_set:
return None
return sorted(current_set | set(turn_end_ids))

270
studio/backend/core/inference/chat_template_helpers.py
View file

@ -3,12 +3,60 @@
"""
Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
fallback for templates that reject reasoning/tools args.
fallback for templates that reject reasoning/tools args, plus the shared
native-chat-template fallback used by the transformers and MLX backends.
"""
import copy
import json
import logging
from typing import Optional
logger = logging.getLogger(__name__)
def _normalize_tool_call_arguments(messages: list) -> list:
"""Coerce each assistant ``tool_calls[].function.arguments`` from a JSON
string to a dict.
The OpenAI wire format carries ``arguments`` as a JSON string, but some chat
templates (e.g. the stricter Qwen tool templates shipped with mlx-community
checkpoints) iterate ``arguments.items()`` and raise
``TypeError: Can only get item pairs from a mapping.`` on the string form
when a prior tool call is re-rendered on the next turn. A dict works on both
strict and lenient templates, so parse the string; leave non-JSON or non-dict
values untouched. Returns the original list unchanged when nothing needed
coercing (no copy)."""
mutated = False
out: list = []
for msg in messages:
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
if not tool_calls:
out.append(msg)
continue
new_calls = []
msg_changed = False
for call in tool_calls:
fn = call.get("function") if isinstance(call, dict) else None
args = fn.get("arguments") if isinstance(fn, dict) else None
if isinstance(args, str):
try:
parsed = json.loads(args)
except (ValueError, TypeError):
parsed = None
if isinstance(parsed, dict):
call = {**call, "function": {**fn, "arguments": parsed}}
msg_changed = True
new_calls.append(call)
if msg_changed:
out.append({**msg, "tool_calls": new_calls})
mutated = True
else:
out.append(msg)
return out if mutated else messages
def apply_chat_template_for_generation(
tokenizer,
messages: list,
@ -38,21 +86,209 @@ def apply_chat_template_for_generation(
attempts.append(dict(reasoning_kwargs))
attempts.append({})
last_exc: Optional[Exception] = None
for kwargs in attempts:
def _render(msgs: list) -> str:
last_exc: Optional[Exception] = None
for kwargs in attempts:
try:
return tokenizer.apply_chat_template(
msgs,
tokenize = False,
add_generation_prompt = True,
**kwargs,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
try:
return _render(messages)
except Exception:
# Strict tool templates reject the JSON-string ``arguments`` form via
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
# Original messages render first, so working templates stay byte-identical.
normalized = _normalize_tool_call_arguments(messages)
if normalized is messages:
raise
return _render(normalized)
def render_native_template(
*,
model_info: dict,
active_model_name: Optional[str],
messages: list,
tools: list,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> Optional[str]:
"""Render ``messages`` + ``tools`` with the model's NATIVE chat template.
Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit
the ``tools`` schema, so a tool-calling turn silently stops advertising tools.
The native template ships in the model repo and carries the family's
tool-calling syntax. It is loaded straight from the repo (bypassing any
override on the live tokenizer) and cached on ``model_info``. Returns the
rendered prompt only if the native template actually emits the tools (render
differs with vs without tools); otherwise ``None``.
``hf_token`` is the token the model was loaded with -- passed to the repo load
so a gated/private model's native template can still be fetched (otherwise the
fallback fails silently and keeps the override prompt that dropped tools).
``trust_remote_code`` is sourced from ``model_info`` (the value the model was
actually loaded with) rather than a call-site argument, so the native-template
reload uses exactly the consent already granted at load. A custom-code tokenizer
repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is
passed, so without this the fallback fails silently and keeps the tool-dropping
prompt for a model the user already consented to run remote code for. For a LoRA
adapter the reload targets the base model, whose remote code was gated and loaded
under the same stored flag, so re-passing it executes no unconsented code.
"""
# ``apply_fn`` lets a backend inject its own render; defaults to the module helper.
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
native_tpl = model_info.get("native_chat_template")
if native_tpl is None:
# A LoRA adapter's native template lives on the base model, not the adapter id.
template_source = model_info.get("base_model") or active_model_name
# Re-use the load-time trust_remote_code so a custom-code tokenizer repo can
# instantiate its class (the stored flag already covers template_source).
trust_remote_code = bool(model_info.get("trust_remote_code", False))
try:
return tokenizer.apply_chat_template(
messages,
tokenize = False,
add_generation_prompt = True,
**kwargs,
from transformers import AutoTokenizer
nt = AutoTokenizer.from_pretrained(
template_source,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
native_tpl = nt.chat_template or False
except Exception as exc:
logger.warning(
"Could not load native chat template for '%s': %s",
template_source,
exc,
)
# A failed fetch is not "no template": leave the sentinel unset so the next
# call retries (caching False would pin the tool-dropping override).
return None
model_info["native_chat_template"] = native_tpl
if not native_tpl:
return None
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
if tokenizer is None:
return None
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
# Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the
# generation lock) races concurrent requests.
try:
render_tokenizer = copy.copy(tokenizer)
render_tokenizer.chat_template = native_tpl
except Exception as exc:
logger.warning(
"Could not clone tokenizer for native-template render of '%s': %s",
active_model_name,
exc,
)
return None
try:
with_tools = apply_fn(
render_tokenizer,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
no_tools = apply_fn(
render_tokenizer,
messages,
tools = None,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
logger.warning(
"Native-template tool render failed for '%s': %s",
active_model_name,
exc,
)
return None
return with_tools if with_tools != no_tools else None
def render_with_native_template_fallback(
*,
formatted_prompt: str,
tokenizer,
model_info: dict,
active_model_name: Optional[str],
messages: list,
tools: Optional[list],
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> str:
"""Return ``formatted_prompt``, swapping in a native-template render when an
override template dropped the ``tools`` schema.
If ``tools`` were requested but the live render is identical with and without
them (detected by comparison, robust against tool names in the system prompt),
re-render with the model's native template. Shared by the transformers and MLX
backends so both advertise tools consistently. ``hf_token`` is forwarded so a
gated/private model's native template can still be fetched."""
if not tools:
return formatted_prompt
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
# Probe whether the live template dropped the schema. A tools-requiring template
# can raise here; on any error keep the valid tools prompt rather than lose it.
try:
probe_no_tools = apply_fn(
tokenizer,
messages,
tools = None,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
logger.warning(
"No-tools probe failed for '%s'; keeping the existing tools prompt: %s",
active_model_name,
exc,
)
return formatted_prompt
if formatted_prompt != probe_no_tools:
return formatted_prompt # template already emits the tools schema
native_prompt = render_native_template(
model_info = model_info,
active_model_name = active_model_name,
messages = messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
apply_fn = apply_fn,
hf_token = hf_token,
)
if native_prompt:
logger.info(
"Override template for '%s' dropped tool schemas; using the model's "
"native template for this tool-calling turn.",
active_model_name,
)
return native_prompt
return formatted_prompt

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

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

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

@ -27,6 +27,11 @@ from utils.hardware import (
from core.inference.audio_codecs import AudioCodecManager
from core.inference.runtime_context import runtime_context_length
from core.inference.message_content import content_to_text
from core.inference.chat_eos import (
chat_eos_repair,
resolve_chat_turn_end_eos_ids_using,
)
from core.inference.presence_penalty import _make_presence_penalty_processor
from io import StringIO
import structlog
from loggers import get_logger
@ -210,6 +215,50 @@ class InferenceBackend:
# API uses -1 to disable top-k; transformers uses 0.
return 0 if top_k < 0 else top_k
def _resolve_chat_eos(self, model_name: str) -> None:
"""Resolve this chat model's assistant-turn-end stop tokens once at load,
cache them in model_info, and repair generation_config so every
``.generate()`` path stops at the turn boundary.
Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
``generation_config.json``, so paths that read ``generation_config`` (the
vision path, tool loops) run past the turn and loop. Turn-end markers are
derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
so base/coder models and harmony templates are left untouched.
"""
info = self.models.get(model_name) or {}
model = info.get("model")
container = info.get("tokenizer")
tokenizer = getattr(container, "tokenizer", container) # unwrap processors
if model is None or tokenizer is None:
return
# Vision models carry the chat_template on the processor, not the inner
# tokenizer. Read markers from whichever has one, but resolve ids on the
# generation tokenizer, else the vision path misses the turn-end token.
template_source = container if getattr(container, "chat_template", None) else tokenizer
try:
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
except Exception as e: # never block a load on eos resolution
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
return
info["chat_turn_end_eos_ids"] = turn_end_ids
gen = getattr(model, "generation_config", None)
if gen is None:
return
repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
if repaired is None:
return
previous = gen.eos_token_id
gen.eos_token_id = repaired
logger.info(
"Repaired generation_config.eos_token_id for %s: %s -> %s",
model_name,
previous,
repaired,
)
def load_model(
self,
config: ModelConfig,
@ -221,6 +270,9 @@ class InferenceBackend:
gpu_ids: Optional[list[int]] = None,
) -> bool:
"""Load any model: base, LoRA adapter, text, or vision."""
# Keep the token so the native-template fallback can fetch a
# gated model's repo template later during generation.
self._hf_token = hf_token
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
if max_seq_length <= 0:
max_seq_length = 2048
@ -231,6 +283,8 @@ class InferenceBackend:
# Already loaded?
if model_name in self.models and self.models[model_name].get("model"):
logger.info(f"Model {model_name} already loaded")
if hf_token:
self.models[model_name]["hf_token"] = hf_token
self.active_model_name = model_name
return True
@ -246,6 +300,14 @@ class InferenceBackend:
)
self.models[model_name] = {
# Per-model token: the native-template fallback must use the
# token this model was loaded with, not whichever loaded last.
"hf_token": hf_token,
# Per-model consent: the native-template reload must re-use the
# exact trust_remote_code this model (and a LoRA's base) was loaded
# with, so a custom-code tokenizer repo can be re-fetched without
# executing any code the user did not already consent to.
"trust_remote_code": trust_remote_code,
"is_vision": config.is_vision,
"is_lora": config.is_lora,
"is_audio": config.is_audio,
@ -496,6 +558,7 @@ class InferenceBackend:
max_seq_length,
)
self._resolve_chat_eos(model_name)
self._load_chat_template_info(model_name)
self.active_model_name = model_name
@ -766,9 +829,11 @@ class InferenceBackend:
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
presence_penalty: float = 0.0,
):
"""Run an agentic tool loop on top of ``generate_chat_response``.
@ -802,6 +867,7 @@ class InferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
initial = list(messages)
@ -815,6 +881,7 @@ class InferenceBackend:
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
nudge_tool_calls = nudge_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
@ -837,12 +904,14 @@ class InferenceBackend:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate response for text or vision models (lock held by background thread).
``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking``
are forwarded into ``apply_chat_template`` so templates that understand them
(Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls.
``presence_penalty`` matches the GGUF sampling path (0 disables it).
"""
yield from self._generate_chat_response_inner(
messages = messages,
@ -859,6 +928,7 @@ class InferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
def _generate_chat_response_inner(
@ -878,6 +948,7 @@ class InferenceBackend:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Inner generation logic, called by generate_chat_response and
generate_with_adapter_control.
@ -917,6 +988,7 @@ class InferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event = cancel_event,
presence_penalty = presence_penalty,
)
return
else:
@ -946,6 +1018,22 @@ class InferenceBackend:
tokenizer,
chat_template = template_name,
)
# The mapper installs the effective template only now, at generate
# time, so re-resolve and UNION into the load-time cache (never
# overwrite). get_chat_template can return a remapped tokenizer
# (turn-end folded onto doc-eos) while generate_stream reads the
# original, so take marker strings from the mapped template but
# resolve their ids on the original.
try:
_gen_tok = model_info.get("tokenizer") or tokenizer
refreshed = resolve_chat_turn_end_eos_ids_using(
getattr(tokenizer, "tokenizer", tokenizer),
getattr(_gen_tok, "tokenizer", _gen_tok),
)
existing = model_info.get("chat_turn_end_eos_ids") or []
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
except Exception as e:
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
else:
logger.info(
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
@ -975,6 +1063,27 @@ class InferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# If tools were requested but the (possibly overridden) template ignored
# them, fall back to the model's native template (shared with MLX).
from core.inference.chat_template_helpers import (
render_with_native_template_fallback,
)
formatted_prompt = render_with_native_template_fallback(
formatted_prompt = formatted_prompt,
tokenizer = tokenizer,
model_info = model_info,
active_model_name = self.active_model_name,
messages = template_messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
apply_fn = self._apply_chat_template_for_generation,
hf_token = model_info.get("hf_token"),
)
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
logger.error(f"Error applying chat template: {e}")
@ -992,6 +1101,7 @@ class InferenceBackend:
repetition_penalty,
cancel_event = cancel_event,
_adapter_state = _adapter_state,
presence_penalty = presence_penalty,
)
def _generate_vision_response(
@ -1006,6 +1116,7 @@ class InferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Handle vision model generation with true token-by-token streaming."""
model_info = self.models[self.active_model_name]
@ -1095,6 +1206,14 @@ class InferenceBackend:
top_k = top_k,
min_p = min_p,
)
# Presence penalty (GGUF parity) for VLM chat.
_vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None
if _vision_input_ids is not None:
_pp = _make_presence_penalty_processor(
presence_penalty, int(_vision_input_ids.shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
err: dict[str, str] = {}
@ -1323,11 +1442,13 @@ class InferenceBackend:
repetition_penalty: float = 1.0,
cancel_event = None,
_adapter_state = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate a streaming text response (text models only).
_adapter_state: if not None, the background thread toggles adapters
before model.generate(), under _generation_lock.
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
"""
if not self.active_model_name:
yield "Error: No active model"
@ -1382,11 +1503,18 @@ class InferenceBackend:
min_p = min_p,
repetition_penalty = repetition_penalty,
do_sample = temperature > 0,
eos_token_id = tokenizer.eos_token_id,
# Resolved once at load (chat_template-derived turn-end tokens).
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
pad_token_id = tokenizer.eos_token_id
if tokenizer.pad_token_id is None
else tokenizer.pad_token_id,
)
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
_pp = _make_presence_penalty_processor(
presence_penalty, int(inputs["input_ids"].shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
if cancel_event is not None:
from transformers.generation.stopping_criteria import (
StoppingCriteria,

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

@ -38,8 +38,35 @@ from core.inference.llama_server_args import (
strip_shadowing_flags,
strip_split_mode_only,
)
# Share strip / signal constants with the multi-format parser so BUFFERING also
# catches Llama-3 / Mistral / Gemma 4 (legacy helper only knew <tool_call> / <function=).
from core.inference.tool_call_parser import (
_GEMMA_BARE_TC_PREFIX_RE,
_GEMMA_BARE_TC_RE,
_TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS,
_balanced_brace_end,
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
_strip_glm_calls,
_strip_mistral_closed_calls,
TOOL_XML_SIGNALS as _SHARED_TOOL_XML_SIGNALS,
RAG_MAX_SEARCHES_PER_TURN,
RAG_SEARCH_CAP_NUDGE,
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
strip_leading_bare_json_call,
strip_llama3_leading_sentinels,
strip_tool_markup as _shared_strip_tool_markup,
)
# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated
# pattern lists, so the GGUF streaming strip stays aligned with the parser.
from core.tool_healing import (
_TOOL_ALL_PATS,
_REHEARSAL_TAIL_STRIP_RE,
_strip_bracket_tag_calls,
apply_tool_strip_patterns,
strip_outside_think,
strip_tool_call_markup,
)
from utils.native_path_leases import child_env_without_native_path_secret
@ -49,10 +76,10 @@ from utils.subprocess_compat import (
)
from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs
from core.inference.tool_call_parser import (
RAG_MAX_SEARCHES_PER_TURN,
RAG_SEARCH_CAP_NUDGE,
TOOL_XML_SIGNALS,
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
MAX_ACT_REPROMPTS as _MAX_REPROMPTS,
REPROMPT_MAX_CHARS as _REPROMPT_MAX_CHARS,
is_short_intent_without_action as _is_short_intent_without_action,
reprompt_to_act_message as _reprompt_to_act_message,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
@ -202,25 +229,8 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
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.
_INTENT_SIGNAL = re.compile(
r"(?i)("
# Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes).
# Excludes "I can"/"I should"/"I want to"/"let's" (common in answers).
# Negative lookahead drops negated forms ("I will not") so a refusal
# doesn't trigger a re-prompt.
r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
r"|"
# Step/plan framing: "First ...", "Step 1:", "Here's my plan"
r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
r"|"
# "Now I" / "Next I" patterns
r"\b(?:now i|next i)\b"
r")"
)
_MAX_REPROMPTS = 1
# Plan-without-action re-prompt state (intent signal, caps, message) now lives
# in tool_call_parser, imported above under its old aliases.
# Default max_tokens to the effective context when known. The floor is high
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
@ -231,7 +241,10 @@ _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
# is exempt because it needs immediate artifact feedback.
_PROVISIONAL_ARGS_MIN_CHARS = 256
_DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min
_REPROMPT_MAX_CHARS = 2000
# Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors
# loop). Structured delta.tool_calls are grammar-bounded by llama-server; text
# parsed from content is not, so one runaway turn could fan out unbounded.
_MAX_TOOL_CALLS_PER_TURN = 8
_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
re.I,
@ -242,9 +255,70 @@ _FINAL_ANSWER_SIGNAL = re.compile(
)
def _is_short_intent_without_action(text: str) -> bool:
stripped = text.strip()
return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None
def _gguf_active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in (active_tools or [])
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
# Rehearsal NAME chars (word + hyphen, matching the parser); the lookbehind excludes the
# Mistral [CALL_ID]...[ARGS] shape.
_GGUF_REHEARSAL_ARGS_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]")
def _gguf_rehearsal_signal_pos(text: str, active_tools: list[dict]) -> int:
"""Index of the first ``NAME[ARGS]`` whose NAME is an active tool, else -1. A
bare/inactive-name ``foo[ARGS]`` in prose is not a call; mirrors the safetensors
``_earliest_tool_signal`` name-gating (no unrestricted GGUF mode)."""
active = set(_gguf_active_tool_names(active_tools))
if not active:
return -1
for m in _GGUF_REHEARSAL_ARGS_RE.finditer(text):
if m.group(1) in active:
return m.start()
return -1
def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict]) -> bool:
"""True when ``text`` holds a genuine tool-call boundary for one of ``signals``.
Unambiguous markers (``<tool_call>``, ``[TOOL_CALLS]``, ``<function=``) count on a
plain substring hit; an ``[ARGS]`` hit is genuine only when an active tool name
precedes it, so inactive-name prose is neither drained nor parsed."""
for sig in signals:
if sig == "[ARGS]":
if _gguf_rehearsal_signal_pos(text, active_tools) >= 0:
return True
continue
if sig in text:
return True
return False
def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool:
"""True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an
active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``.
Mirrors the safetensors loop so the split rehearsal call is not streamed."""
if not stripped or any(ch.isspace() for ch in stripped):
return False
for name in _gguf_active_tool_names(active_tools):
if stripped == name or f"{name}[ARGS]".startswith(stripped):
return True
return False
def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int:
"""Length of a trailing bare tool-name token that may be a split rehearsal call
(``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it
instead of leaking the name. Returns 0 for ordinary prose. Mirrors safetensors."""
i = len(text)
while i > 0 and not text[i - 1].isspace():
i -= 1
tail = text[i:]
return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0
def _should_suppress_forced_no_tool_output(text: str) -> bool:
@ -545,6 +619,13 @@ _TOOL_TEMPLATE_MARKERS = (
"'role' == 'tool'",
'message.role == "tool"',
"message.role == 'tool'",
# DeepSeek: no top-level ``{% if tools %}`` block; it gates emission on
# ``message['role'] == 'tool'`` plus ``message['tool_calls'] is defined``.
"message['role'] == 'tool'",
'message["role"] == "tool"',
"message['tool_calls']",
'message["tool_calls"]',
"tool_calls is defined",
)
@ -605,6 +686,16 @@ def detect_reasoning_flags(
else []
)
if effort_levels:
# DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its
# template only branches on 'max', so the literal scan misses 'high'. Add it
# (matched on whole repo-name segments, so 'deepseek-v40' won't false-match)
# to expose the full none/high/max ladder instead of none/max.
segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1])
is_dsv4 = "deepseek4" in segments or any(
a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:])
)
if is_dsv4 and "high" not in effort_levels:
effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index)
# GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort
# level among a discrete set (e.g. 'high' | 'max'). Distinct from
# gpt-oss (reasoning_effort only, no on/off gate) and Qwen
@ -1660,9 +1751,13 @@ class LlamaCppBackend:
# 'low' effort the way gpt-oss does (those models genuinely
# cannot disable).
thinking_off = enable_thinking is False or reasoning_effort == "none"
if enable_thinking is not None or reasoning_effort == "none":
# A named effort level implies thinking on, so emit enable_thinking
# even if the caller sent only reasoning_effort (else the template
# defaults it off and the requested level never renders).
effort_on = reasoning_effort in self._reasoning_effort_levels
if enable_thinking is not None or reasoning_effort == "none" or effort_on:
kwargs["enable_thinking"] = not thinking_off
if not thinking_off and reasoning_effort in self._reasoning_effort_levels:
if not thinking_off and effort_on:
kwargs["reasoning_effort"] = reasoning_effort
elif self._reasoning_style == "reasoning_effort":
if reasoning_effort in ("none", "low", "medium", "high"):
@ -3048,6 +3143,12 @@ class LlamaCppBackend:
_CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch)
_CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x)
_CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok)
# DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large
# context-scaling compute buffer the rates above miss (present even with an f16
# cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without
# it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU.
_DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch
_DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M)
def _estimate_compute_buffer_bytes(
self,
@ -3097,6 +3198,14 @@ class LlamaCppBackend:
if n_embd <= 0 or n_ctx <= 0:
return 0
ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
if getattr(self, "_architecture", None) == "deepseek4":
# DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires
# for any KV type -- the indexer scratch is present even with an f16 cache.
ub_scale = ub / self._DEFAULT_N_UBATCH
return int(
self._DSV4_CTX_COMPUTE_FLAT_BYTES
+ self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale
)
if _kv_bytes_per_elem(cache_type_kv) < 2.0:
# Quantized cache: the dequant scratch dominates and scales with n_embd.
# MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on
@ -7881,12 +7990,17 @@ class LlamaCppBackend:
# ── Message building (OpenAI format) ──────────────────────────
@staticmethod
def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]:
"""Thin wrapper around the shared parser in tool_call_parser
so safetensors and llama_cpp pick up the same fixes."""
def _parse_tool_calls_from_text(
content: str,
*,
allow_incomplete: bool = True,
enabled_tool_names: Optional[set] = None,
) -> list[dict]:
"""Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form."""
return _shared_parse_tool_calls_from_text(
content,
allow_incomplete = allow_incomplete,
enabled_tool_names = enabled_tool_names,
)
@staticmethod
@ -8353,6 +8467,7 @@ class LlamaCppBackend:
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
@ -8392,12 +8507,23 @@ class LlamaCppBackend:
_reasoning_started_at: Optional[float] = None
_reasoning_summary_emitted = False
# Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the
# ORIGINAL tools list so a spent one-shot still reads as a tool name. None = no gate.
_enabled_names_gate = set(_gguf_active_tool_names(tools)) if tools else None
# Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent
# one-shot), else its repeat is stripped but never drained and the turn ends blank.
_detect_tools = list(tools or [])
def _reasoning_summary_event(started_at: float) -> dict:
return {
"type": "reasoning_summary",
"duration_ms": round((time.monotonic() - started_at) * 1000.0),
}
# Enabled-name gate for the markerless Gemma strip (disabled/example
# names stay visible). Set per iteration; None = pre-loop name-agnostic.
_enabled_tool_names = None
def _strip_tool_markup(
text: str,
*,
@ -8406,14 +8532,42 @@ class LlamaCppBackend:
) -> str:
if not (auto_heal_tool_calls or force):
return text
return strip_tool_call_markup(text, final = final)
# Delegate to the shared parser-side strip so the GGUF cleanup covers every family the
# parser promotes (Llama <|python_tag|>, Mistral [TOOL_CALLS], bare rehearsal, function
# XML, Gemma) and stays aligned with detection; tool_healing's strip omits the loop-only
# forms (python_tag / Mistral name) and would leak them into display.
return _shared_strip_tool_markup(
text, final = final, enabled_tool_names = _enabled_names_gate
)
def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
if not (auto_heal_tool_calls or force):
return text
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
def _seg(segment: str, is_last: bool) -> str:
# Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced
# strips first (nested JSON removed whole; literal markup inside a value is that
# call's data), then the guarded function-XML / GLM scans, then the regex arms
# (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last
# segment (a bare ``foo[ARGS]`` before <think> is prose). Rehearsal + markerless
# strips are name-gated on the ORIGINAL list (strip/detect aligned).
seg = _strip_mistral_closed_calls(segment)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate)
seg = _strip_function_xml_calls(seg, final = is_last)
seg = _strip_glm_calls(seg, final = is_last)
pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS
for pat in pats:
seg = pat.sub("", seg)
if is_last:
seg = apply_tool_strip_patterns(
seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate
)
return seg
# Preserve think blocks verbatim (a rehearsed call inside one must not be deleted).
return strip_outside_think(text, _seg)
def _build_metadata_event(usage, timings, finish_reason):
"""Final usage+timings metadata event for the given pass, merging its
@ -8449,13 +8603,38 @@ class LlamaCppBackend:
}
def _flush_reasoning_and_buffer():
"""Append buffered reasoning (as a <think> block) then the held
"""Close a live-streamed <think> block (or emit the buffered reasoning
as one block if it never streamed), then append the held
content_buffer to the cumulative display text."""
nonlocal cumulative_display
if reasoning_accum:
nonlocal cumulative_display, in_thinking
if in_thinking:
cumulative_display += "</think>"
in_thinking = False
elif reasoning_accum:
cumulative_display += "<think>" + reasoning_accum + "</think>"
cumulative_display += content_buffer
def _close_streamed_think() -> bool:
"""Close a live-streamed <think> before a tool call drains, so
consumers without a reasoning extractor (Anthropic) get a balanced
block. Returns True when the caller should yield the result."""
nonlocal cumulative_display, in_thinking, _last_emitted
if not in_thinking:
return False
cumulative_display += "</think>"
in_thinking = False
if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output:
_last_emitted = cumulative_display
return True
return False
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool:
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
probe = strip_llama3_leading_sentinels(text.lstrip())
if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
return False
return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
tool_controller = ToolLoopController(
tools = tools,
auto_heal_tool_calls = auto_heal_tool_calls,
@ -8469,18 +8648,21 @@ class LlamaCppBackend:
)
_MAX_BUFFER_CHARS = 32
# Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors).
_MAX_BARE_JSON_BUFFER = 16384
_append_budget_exhausted_nudge = True
# RAG: cap knowledge-base searches per assistant turn. The controller is
# tool-agnostic, so this gate stays in the loop.
_kb_search_count = 0
# ── Re-prompt on plan-without-action ─────────────────
# When the model describes what it intends to do (forward-looking
# language) without calling a tool, re-prompt once. Only triggers on
# responses signaling intent/planning -- a direct answer like "4" or
# "Hello!" won't match. Pattern compiled at module level
# (_INTENT_SIGNAL).
# Model describes intent without calling a tool: re-prompt once. A
# direct answer ("4", "Hello!") won't match. Pattern shared with the
# safetensors loop (tool_call_parser.INTENT_SIGNAL).
_reprompt_count = 0
# Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved
# re-prompt slots don't extend the budget. Mirrors the safetensors guard.
_tool_iters_done = 0
_forced_tool_call_pending = False
# Reserve extra iterations for re-prompts so they don't consume the
@ -8489,12 +8671,21 @@ class LlamaCppBackend:
for iteration in range(max_tool_iterations + _extra):
if cancel_event is not None and cancel_event.is_set():
return
# Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
_turn_executed_real_tool = False
active_tools = tool_controller.active_tools()
if not active_tools:
_append_budget_exhausted_nudge = False
break
_tool_xml_signals = TOOL_XML_SIGNALS
# Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
_enabled_tool_names = {
(tool.get("function") or {}).get("name")
for tool in active_tools
if (tool.get("function") or {}).get("name")
}
# Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows (like safetensors).
_tool_xml_signals = _SHARED_TOOL_XML_SIGNALS
# Build payload -- stream: True so we detect tool signals
# in the first 1-2 chunks without a non-streaming penalty.
@ -8624,6 +8815,10 @@ class LlamaCppBackend:
# the structured tool call.
has_structured_tc = True
detect_state = _S_DRAINING
# Close the reasoning prefix before the tool card
# (mirrors the is_match path).
if _close_streamed_think():
yield {"type": "content", "text": cumulative_display}
for tc_d in tc_deltas:
idx = tc_d.get("index", 0)
if idx not in tool_calls_acc:
@ -8709,17 +8904,17 @@ class LlamaCppBackend:
continue
# ── Reasoning tokens ──
# Yield only in STREAMING. In BUFFERING and
# DRAINING, accumulate silently so we don't
# corrupt the consumer's prev_text tracker
# (routes/inference.py never resets it
# between tool iterations).
# Stream live except while DRAINING: reasoning is
# orthogonal to tool detection (content_buffer
# only), and the route resets prev_text on
# tool_start, so the <think> block stays a
# monotonic prefix like the no-tool path.
reasoning = delta.get("reasoning_content", "")
if reasoning:
if _reasoning_started_at is None:
_reasoning_started_at = time.monotonic()
reasoning_accum += reasoning
if detect_state == _S_STREAMING:
if detect_state != _S_DRAINING:
if not in_thinking:
cumulative_display += "<think>"
in_thinking = True
@ -8752,12 +8947,18 @@ class LlamaCppBackend:
in_thinking = False
cumulative_display += token
cleaned = _strip_tool_markup_streaming(cumulative_display)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
# Hold a trailing bare active-tool-name (split rehearsal)
# until [ARGS] arrives; released by later prose or stream end.
_hold = _held_rehearsal_tail_len(cleaned, _detect_tools)
_emit = (
cleaned[: len(cleaned) - _hold] if _hold else cleaned
)
if len(_emit) > len(_last_emitted):
_last_emitted = _emit
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
"text": _emit,
}
elif detect_state == _S_BUFFERING:
@ -8766,7 +8967,8 @@ class LlamaCppBackend:
if not stripped_buf:
continue
# Check tool signal prefixes.
# Bracket tags arrive mid-buffer, so substring-check too;
# ``[ARGS]`` counts only as a regex-matched NAME[ARGS].
is_prefix = False
is_match = False
for sig in _tool_xml_signals:
@ -8776,14 +8978,91 @@ class LlamaCppBackend:
if sig.startswith(stripped_buf):
is_prefix = True
break
if sig == "[ARGS]":
# Active NAME[ARGS] only; inactive-name prose
# is gated out, not drained/parsed.
if (
_gguf_rehearsal_signal_pos(
stripped_buf, _detect_tools
)
>= 0
):
is_match = True
break
elif sig.startswith("[") and sig in stripped_buf:
is_match = True
break
if is_match:
# Split rehearsal: hold the bare name until
# its [ARGS] arrives and matches above.
is_rehearsal_prefix = False
if (
not is_match
and not is_prefix
and _is_rehearsal_prefix(stripped_buf, _detect_tools)
):
is_prefix = True
is_rehearsal_prefix = True
# Signal-less call shapes (mirror the safetensors
# loop): Llama-3.2 bare {"name":..} and Gemma
# call:NAME{...} would otherwise stream raw.
_hold_buffer = False
# Whole buffer is the call (no visible prefix) -- drain silently.
_drain_silently = False
if not is_match and not is_prefix:
_bare = strip_llama3_leading_sentinels(stripped_buf)
if _bare.startswith("{"):
if _balanced_brace_end(_bare, 0) is None:
if len(stripped_buf) < _MAX_BARE_JSON_BUFFER:
_hold_buffer = True
elif _looks_like_enabled_bare_json(
_bare, _enabled_tool_names
):
# Oversized still-open enabled call: drain
# rather than leak; a giant ordinary JSON
# answer still streams.
_drain_silently = True
elif self._parse_tool_calls_from_text(
content_buffer,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
):
_drain_silently = True
elif (
"call:".startswith(stripped_buf)
or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf)
is not None
or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None
):
# Whitespace-tolerant like the parser.
if _GEMMA_BARE_TC_RE.match(stripped_buf):
_drain_silently = True
elif len(stripped_buf) < _MAX_BUFFER_CHARS:
_hold_buffer = True
if _drain_silently:
# The buffered content IS the call; drain it
# without yielding. A live <think> prefix is
# separate from it -- close that.
detect_state = _S_DRAINING
if _close_streamed_think():
yield {
"type": "content",
"text": cumulative_display,
}
elif is_match:
# Tool signal -- flush any visible
# prefix before DRAINING so the
# route sends it before tool_start.
# Use the final strip (all families incl. Llama
# <|python_tag|> / Mistral name): the buffer holds
# the whole call, so a streaming closed-only strip
# would leak its open-ended markup as display text.
_flush_reasoning_and_buffer()
cleaned = _strip_tool_markup_streaming(
cleaned = _strip_tool_markup(
cumulative_display,
final = True,
force = True,
)
if len(cleaned) > len(_last_emitted):
@ -8794,7 +9073,15 @@ class LlamaCppBackend:
"text": cleaned,
}
detect_state = _S_DRAINING
elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS:
elif _hold_buffer or (
is_prefix
and (
is_rehearsal_prefix
or len(stripped_buf) < _MAX_BUFFER_CHARS
)
):
# A rehearsal prefix is self-bounded; the buffer
# cap must not cut long MCP names short.
pass # keep buffering
else:
# Not a tool -- flush buffer
@ -8805,12 +9092,20 @@ class LlamaCppBackend:
cleaned = _strip_tool_markup(
cumulative_display,
)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
# Same trailing-name hold as STREAMING for this
# first flush out of BUFFERING.
_hold = _held_rehearsal_tail_len(cleaned, _detect_tools)
_emit = (
cleaned[: len(cleaned) - _hold]
if _hold
else cleaned
)
if len(_emit) > len(_last_emitted):
_last_emitted = _emit
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
"text": _emit,
}
except json.JSONDecodeError:
@ -8821,7 +9116,18 @@ class LlamaCppBackend:
# ── Resolve BUFFERING at stream end ──
if detect_state == _S_BUFFERING:
stripped_buf = content_buffer.lstrip()
if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals):
# A held bare-JSON fragment has no XML signal; route it to DRAINING (the signal-only
# gate below would flush the raw JSON to the user).
_bare_eos = strip_llama3_leading_sentinels(stripped_buf)
# Gate on enabled names so an ordinary JSON answer isn't routed to DRAINING and dropped.
_is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json(
_bare_eos, _enabled_tool_names
)
if stripped_buf and _gguf_has_genuine_tool_signal(
stripped_buf, _tool_xml_signals, _detect_tools
):
detect_state = _S_DRAINING
elif _is_bare_tc:
detect_state = _S_DRAINING
elif content_accum or reasoning_accum:
detect_state = _S_STREAMING
@ -8837,7 +9143,9 @@ class LlamaCppBackend:
),
}
elif reasoning_accum and not has_content_tokens:
# Reasoning-only reply: show it as plain text.
# Reasoning-only reply: show it as the main response,
# not a thinking block (mirrors the no-tool path; the
# route's extractor closes the streamed <think>).
if _reasoning_started_at is not None and not _reasoning_summary_emitted:
_reasoning_summary_emitted = True
yield _reasoning_summary_event(_reasoning_started_at)
@ -8848,20 +9156,26 @@ class LlamaCppBackend:
"text": cumulative_display,
}
else:
# Held buffer was no tool signal and no enabled bare-JSON call: a leading ``{`` is an
# ordinary JSON answer and must be shown; any other partial-markup prefix is dropped.
_held = strip_llama3_leading_sentinels(content_buffer.lstrip())
if _held.startswith("{") and not _suppress_visible_output:
yield {"type": "content", "text": _held}
return
# ── STREAMING path: no tool call ──
if detect_state == _S_STREAMING:
# Safety net: check for XML tool signals in content. The
# Safety net: re-parse the full content for tool calls. The
# route layer resets prev_text on tool_start, so post-tool
# synthesis streams correctly even if content was emitted
# before the tool XML.
_safety_tc = None
if any(s in content_accum for s in _tool_xml_signals):
_safety_tc = self._parse_tool_calls_from_text(
content_accum,
allow_incomplete = auto_heal_tool_calls,
)
# Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less
# calls carry no XML signal, so a signal gate would let them slip past.
_safety_tc = self._parse_tool_calls_from_text(
content_accum,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
)
if not _safety_tc:
# ── Re-prompt on plan-without-action ──
# If the model described its intent (forward-looking
@ -8879,8 +9193,10 @@ class LlamaCppBackend:
r"(?i)\brender[_\s-]?html\b",
_stripped,
)
# None keeps the default-on re-prompt; False disables it.
if (
auto_heal_tool_calls
and (nudge_tool_calls is None or nudge_tool_calls)
and active_tools
and not _render_html_already_done_intent
and _reprompt_count < _MAX_REPROMPTS
@ -8909,12 +9225,7 @@ class LlamaCppBackend:
conversation.append(
{
"role": "user",
"content": (
"You have access to enabled tools. If a tool is needed to satisfy "
"the user's request or complete the action you described, call "
f"{tool_hint} now. If no tool is needed, provide the final answer "
"and follow the user's requested format."
),
"content": _reprompt_to_act_message(tool_hint),
}
)
# Accumulate tokens and timing from this iteration.
@ -8946,6 +9257,12 @@ class LlamaCppBackend:
"type": "content",
"text": forced_visible_text,
}
elif not _suppress_visible_output:
# Turn ended as a plain answer (no [ARGS] followed): the held
# rehearsal tail is real prose, release it.
_final_clean = _strip_tool_markup_streaming(cumulative_display)
if len(_final_clean) > len(_last_emitted):
yield {"type": "content", "text": _final_clean}
# Content was already streamed. Yield metadata.
yield {"type": "status", "text": ""}
@ -8978,10 +9295,13 @@ class LlamaCppBackend:
for i in sorted(tool_calls_acc)
if (tool_calls_acc[i].get("function", {}).get("name", "").strip())
] or None
if not tool_calls and any(s in content_accum for s in _tool_xml_signals):
if not tool_calls:
# Unconditional re-parse: we only reach DRAINING when the buffer looked like a
# call, and bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on.
tool_calls = self._parse_tool_calls_from_text(
content_accum,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
)
if tool_calls and not has_structured_tc:
content_text = _strip_tool_markup(
@ -8989,6 +9309,11 @@ class LlamaCppBackend:
final = True,
force = True,
)
# ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call so the
# executed call isn't replayed as text or next-turn history.
content_text = strip_leading_bare_json_call(
content_text, _enabled_tool_names
)
if tool_calls:
logger.info(
f"Parsed {len(tool_calls)} tool call(s) from "
@ -9002,6 +9327,13 @@ class LlamaCppBackend:
if content_accum:
# Strip leaked tool-call XML before yielding.
content_accum = _strip_tool_markup(content_accum, final = True)
# A truncated bare-JSON call has no XML markup to strip and didn't parse. With
# Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched);
# off keeps it visible per the strict contract.
if content_accum and active_tools and auto_heal_tool_calls:
content_accum = strip_leading_bare_json_call(
content_accum, _enabled_tool_names
)
if content_accum:
yield {"type": "content", "text": content_accum}
_meta = _build_metadata_event(
@ -9019,6 +9351,29 @@ class LlamaCppBackend:
_accumulated_predicted_ms += _it.get("predicted_ms", 0)
_accumulated_predicted_n += _it.get("predicted_n", 0)
# Collapse exact-duplicate calls and cap the count for the TEXTUAL
# fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN).
if tool_calls and not has_structured_tc and len(tool_calls) > 1:
_seen_keys: set = set()
_deduped: list = []
for _tc in tool_calls:
_fn = _tc.get("function", {}) or {}
_key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
if _key in _seen_keys:
continue
_seen_keys.add(_key)
_deduped.append(_tc)
if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN:
break
if len(_deduped) != len(tool_calls):
logger.info(
"GGUF textual fallback: collapsed %d repeated tool call(s) "
"in one turn to %d",
len(tool_calls),
len(_deduped),
)
tool_calls = _deduped
# disable_parallel_tool_use: execute only the first tool call
# this turn. Truncate before building assistant_msg so the
# conversation stays consistent and extra calls are never executed.
@ -9144,6 +9499,8 @@ class LlamaCppBackend:
_kb_search_count += 1
completion = tool_controller.record_result(decision, result)
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
# A tool ran this turn, so it counts against the caller's budget.
_turn_executed_real_tool = True
yield completion.tool_end_event()
conversation.append(completion.tool_message())
@ -9167,6 +9524,12 @@ class LlamaCppBackend:
if tool_controller.force_final_answer or not tool_controller.active_tools():
_append_budget_exhausted_nudge = False
break
# Count only real tool turns against the cap so reserved re-prompt slots can't become
# extra tool rounds; a no-op correction turn doesn't consume budget (GGUF parity).
if _turn_executed_real_tool:
_tool_iters_done += 1
if _tool_iters_done >= max_tool_iterations:
break
continue
except httpx.ConnectError:

112
studio/backend/core/inference/mlx_inference.py
View file

@ -41,6 +41,50 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
}
def _make_mlx_presence_penalty_processor(penalty: float):
"""Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path.
generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the
full running sequence; the first call is prompt-only, so latch that length
and penalize only after it.
"""
state = {"prompt_len": None}
def _processor(tokens, logits):
if state["prompt_len"] is None:
# First call = prompt only; latch its length.
state["prompt_len"] = int(tokens.shape[0])
return logits
generated = tokens[state["prompt_len"] :]
if generated.size == 0:
return logits
import mlx.core as mx
vocab = logits.shape[-1]
# Bound generated ids to the valid range [0, vocab) before they index
# logits. MLX does no bounds checking and out-of-bounds indexing is
# documented undefined behavior (crash / memory corruption), unlike the
# torch path's harmless negative wrap -- so this bound is load-bearing
# here and matches the torch filter seen[(seen >= 0) & (seen < vocab)].
# MLX has no boolean-mask filtering (data-dependent output shape is
# unsupported), so instead of compacting the id list we route every
# out-of-range or negative id to a scratch slot at index ``vocab`` that
# is dropped before the subtract. That scratch slot can never collide
# with a real token, so real ids (including id 0) are penalized exactly
# once and stray ids are ignored.
valid = (generated >= 0) & (generated < vocab)
safe = mx.where(valid, generated, vocab).astype(mx.int32)
# Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate
# ids are idempotent, so presence applies once per distinct token; the
# scratch column is discarded and the full-width subtract stays on-device.
mask = mx.zeros((vocab + 1,), dtype = logits.dtype)
mask[safe] = penalty
logits = logits - mask[:vocab]
return logits
return _processor
class MLXInferenceBackend:
def __init__(self):
self.models = {}
@ -104,6 +148,9 @@ class MLXInferenceBackend:
) -> bool:
import mlx.core as mx
# Keep the token so the native-template fallback can fetch a
# gated model's repo template later during generation.
self._hf_token = hf_token
model_name = config.identifier if hasattr(config, "identifier") else str(config)
is_vision = getattr(config, "is_vision", False)
@ -168,11 +215,20 @@ class MLXInferenceBackend:
self.active_model_name = model_name
self.models[model_name] = {
# Per-model token for the native-template fallback (matches transformers).
"hf_token": hf_token,
# Per-model consent for the native-template reload: re-use the exact
# trust_remote_code this model was loaded with (matches transformers).
"trust_remote_code": trust_remote_code,
"model": self._model,
"tokenizer": self._tokenizer,
"processor": self._processor,
"is_vision": is_vision,
"is_lora": getattr(config, "is_lora", False),
# For a LoRA adapter the native chat template lives on the base model.
"base_model": getattr(config, "base_model", None)
if getattr(config, "is_lora", False)
else None,
"is_audio": False,
"audio_type": None,
"has_audio_input": False,
@ -270,6 +326,7 @@ class MLXInferenceBackend:
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
) -> Generator[str, None, None]:
if self._model is None:
raise RuntimeError("No model loaded")
@ -317,6 +374,7 @@ class MLXInferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
else:
yield from self._generate_text(
@ -332,6 +390,7 @@ class MLXInferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
def _generate_text(
@ -349,12 +408,14 @@ class MLXInferenceBackend:
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
):
from mlx_lm import stream_generate
from mlx_lm.sample_utils import make_sampler, make_logits_processors
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
render_with_native_template_fallback,
)
prompt = apply_chat_template_for_generation(
@ -368,6 +429,25 @@ class MLXInferenceBackend:
if prompt is None:
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
# Same parity fix as the transformers backend: if the template dropped the
# requested tools, fall back to the native template so MLX text models keep
# advertising them. ``self._tokenizer`` is this entry's model_info tokenizer,
# so probe and native render share a renderer. (The VLM path renders via the
# processor for image tokens and is intentionally not wired here.)
model_info = self.models.get(self.active_model_name, {})
prompt = render_with_native_template_fallback(
formatted_prompt = prompt,
tokenizer = self._tokenizer,
model_info = model_info,
active_model_name = self.active_model_name,
messages = messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
hf_token = model_info.get("hf_token"),
)
sampler = make_sampler(
temp = temperature,
top_p = top_p,
@ -375,15 +455,21 @@ class MLXInferenceBackend:
min_p = float(min_p or 0.0),
min_tokens_to_keep = 1,
)
# Only build a logits processor for a non-trivial repetition penalty.
logits_processors = None
# Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths).
logits_processors = []
if repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
1.0,
):
logits_processors = make_logits_processors(
repetition_penalty = float(repetition_penalty),
logits_processors.extend(
make_logits_processors(
repetition_penalty = float(repetition_penalty),
)
)
if presence_penalty:
logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
if not logits_processors:
logits_processors = None
token_ids = []
logger.info(
@ -449,6 +535,7 @@ class MLXInferenceBackend:
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
):
from mlx_vlm import stream_generate as vlm_stream
@ -496,10 +583,23 @@ class MLXInferenceBackend:
top_k = int(top_k or 0),
min_p = float(min_p or 0.0),
)
if repetition_penalty is not None and float(repetition_penalty) not in (
_rep_active = repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
1.0,
):
)
if presence_penalty:
# Presence needs a custom processor: pass the full list (repetition +
# presence) instead of the repetition_penalty shortcut so both apply once.
from mlx_lm.sample_utils import make_logits_processors
_vlm_processors = []
if _rep_active:
_vlm_processors.extend(
make_logits_processors(repetition_penalty = float(repetition_penalty))
)
_vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
vlm_kwargs["logits_processors"] = _vlm_processors
elif _rep_active:
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
with self._generation_lock:

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

@ -45,6 +45,10 @@ _DISPATCH_STOP_TIMEOUT = 5.0
_DISPATCH_IDLE_TIMEOUT = 30.0
_DISPATCH_DRAIN_TIMEOUT = 5.0
# Max wait for a cancelled generation to release _gen_lock before unload_model
# tears the subprocess down. Only bounds a wedged worker.
_UNLOAD_GEN_LOCK_TIMEOUT = 15.0
class InferenceOrchestrator:
"""
@ -60,7 +64,13 @@ class InferenceOrchestrator:
self._cmd_queue: Any = None
self._resp_queue: Any = None
self._cancel_event: Any = None # mp.Event — set to cancel generation
# Set for the whole unload; the worker never clears it (unlike _cancel_event),
# so a generate queued behind the cancelled one is skipped, not run.
self._drain_event: Any = None
self._gen_lock = threading.Lock() # Serializes generation
# Set during a switch so a generation winning the _gen_lock handoff bails
# instead of starting on the outgoing model.
self._unload_pending = False
# Dispatcher state for compare mode (adapter-controlled requests):
# bypass _gen_lock, send commands directly, read from per-request
@ -69,6 +79,12 @@ class InferenceOrchestrator:
self._mailbox_lock = threading.Lock()
self._dispatcher_thread: Optional[threading.Thread] = None
self._dispatcher_stop = threading.Event()
# Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses
# _gen_lock, so two concurrent compare requests can both reach _start_dispatcher;
# without this lock both could observe no live dispatcher and each spawn one,
# orphaning the extra thread (self._dispatcher_thread tracks only the last). The
# orphan later steals the "unloaded" reply off resp_queue and hangs unload_model.
self._dispatcher_lifecycle_lock = threading.Lock()
# Local state mirrors (updated from subprocess responses)
self.active_model_name: Optional[str] = None
@ -92,13 +108,11 @@ class InferenceOrchestrator:
@property
def default_models(self) -> list[str]:
# Wait up to 5s for background HF fetch
self._top_models_ready.wait(timeout = 5)
top_gguf = self._top_gguf_cache or []
top_hub = self._top_hub_cache or []
# Curated static defaults first, then HF download-ranked to backfill.
# Send extras so the frontend keeps 4 per category after removing
# downloaded ones.
# Never wait for the remote Hugging Face ranking during startup. Chat's
# first /api/models/list needs curated defaults immediately; the
# background fetch backfills extra choices on later calls.
result: list[str] = []
seen: set[str] = set()
for m in self._static_models + top_gguf + top_hub:
@ -159,6 +173,7 @@ class InferenceOrchestrator:
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._cancel_event = _CTX.Event()
self._drain_event = _CTX.Event()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
@ -167,6 +182,7 @@ class InferenceOrchestrator:
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"cancel_event": self._cancel_event,
"drain_event": self._drain_event,
"config": config,
},
daemon = True,
@ -228,6 +244,7 @@ class InferenceOrchestrator:
self._cmd_queue = None
self._resp_queue = None
self._cancel_event = None
self._drain_event = None
logger.info("Inference subprocess shut down")
def _cleanup(self):
@ -409,6 +426,7 @@ class InferenceOrchestrator:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> dict:
"""Build the 'generate' command shared by the locked and dispatched paths."""
cmd = {
@ -423,6 +441,7 @@ class InferenceOrchestrator:
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
"presence_penalty": presence_penalty,
}
# Only forward template kwargs the caller set, for older worker compat.
if use_adapter is not None:
@ -456,7 +475,15 @@ class InferenceOrchestrator:
cancel ack from that same source so stale events don't leak into the
next request.
"""
# Latch this stream's subprocess/queue: if a wedged worker is torn down and a
# later load spawns a fresh one, bail rather than re-block on the new queue
# under _gen_lock (deadlock).
initial_proc = self._proc
initial_resp_queue = self._resp_queue
while True:
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
yield f"Error: {self._subprocess_crash_message(crash_context)}"
return
resp = read_one(read_timeout)
if resp is None:
# Check subprocess health
@ -493,33 +520,56 @@ class InferenceOrchestrator:
# Dispatcher — per-request mailbox routing for compare mode
# ------------------------------------------------------------------
def _start_dispatcher(self) -> None:
def _start_dispatcher(self) -> bool:
"""Start the dispatcher thread if not already running.
The dispatcher reads the shared resp_queue and routes responses to
per-request mailbox queues, letting multiple adapter-controlled
(compare) requests be in-flight without holding _gen_lock.
"""
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
return
self._dispatcher_stop.clear()
self._dispatcher_thread = threading.Thread(
target = self._dispatcher_loop,
daemon = True,
name = "inference-dispatcher",
)
self._dispatcher_thread.start()
logger.debug("Dispatcher thread started")
The whole check-then-spawn runs under _dispatcher_lifecycle_lock so
concurrent compare requests (which bypass _gen_lock) can't both observe
no live dispatcher and each spawn one. Returns True only for the caller
that actually started a new thread; False if one was already alive.
"""
with self._dispatcher_lifecycle_lock:
# Refuse to start while an unload is in progress. unload_model sets
# _unload_pending under this same lock before it stops the idle
# dispatcher, so a start queued behind that stop observes the unload
# here and bails. Without this a fresh dispatcher would be spawned
# after the stop, become the resp_queue reader, and consume the
# worker's "unloaded" reply (unroutable, so dropped) before
# unload_model's _wait_response sees it -- hanging the unload 300s.
if self._unload_pending:
return False
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
return False
self._dispatcher_stop.clear()
self._dispatcher_thread = threading.Thread(
target = self._dispatcher_loop,
daemon = True,
name = "inference-dispatcher",
)
self._dispatcher_thread.start()
logger.debug("Dispatcher thread started")
return True
def _stop_dispatcher(self) -> None:
"""Signal the dispatcher to stop and wait for it."""
if self._dispatcher_thread is None:
return
self._dispatcher_stop.set()
self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
self._dispatcher_thread = None
logger.debug("Dispatcher thread stopped")
"""Signal the dispatcher to stop and wait for it.
Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so
a stop can't interleave with a concurrent start. Callers must NOT hold
_mailbox_lock here: this joins the dispatcher, and the dispatcher loop
takes _mailbox_lock, so holding it would deadlock the join.
"""
with self._dispatcher_lifecycle_lock:
if self._dispatcher_thread is None:
return
self._dispatcher_stop.set()
self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
self._dispatcher_thread = None
logger.debug("Dispatcher thread stopped")
def _dispatcher_loop(self) -> None:
"""Background loop: read resp_queue → route to mailboxes by request_id."""
@ -581,6 +631,7 @@ class InferenceOrchestrator:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Dispatched generation — sends command without holding _gen_lock.
@ -595,9 +646,26 @@ class InferenceOrchestrator:
if not self.active_model_name:
yield "Error: No active model"
return
# Latch the target model so the recheck below can detect a switch that completed
# between _start_dispatcher and mailbox registration (mirrors the locked path's
# expected_model check).
expected_model = self.active_model_name
# Ensure dispatcher is running
self._start_dispatcher()
# Switch in flight (unload waiting on _gen_lock). This path bypasses the lock,
# so without this early-out a compare request would enqueue a generate on the
# outgoing model and delay the switch.
if self._unload_pending:
yield "Error: model is being unloaded"
return
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
# _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned
# the thread, so at most one dispatcher ever exists even when two compare requests race
# here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked
# is_alive() read): if THIS call started the dispatcher and then bails on a racing
# unload, it must stop it again (see the unloading bail below).
started = self._start_dispatcher()
dispatcher_preexisting = not started
request_id = str(uuid.uuid4())
@ -617,6 +685,7 @@ class InferenceOrchestrator:
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
@ -624,10 +693,42 @@ class InferenceOrchestrator:
preserve_thinking = preserve_thinking,
)
# Create mailbox BEFORE sending command
# Create the mailbox BEFORE sending, rechecking _unload_pending under
# _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle
# reads _mailboxes under the same lock, so either the idle check sees this
# mailbox (and tears the dispatcher down) or we see the unload and bail.
# Registering after would orphan the mailbox and hang the compare stream forever.
mailbox: queue.Queue = queue.Queue()
with self._mailbox_lock:
self._mailboxes[request_id] = mailbox
# _unload_pending alone is not enough: an unload that ran fully since
# _start_dispatcher clears it in its finally and stops the dispatcher, so it
# reads False here though the dispatcher is gone and the model swapped. Also
# bail when the active model changed or the dispatcher died: a mailbox with no
# dispatcher to route gen_done/gen_error hangs the compare stream.
dispatcher_alive = (
self._dispatcher_thread is not None and self._dispatcher_thread.is_alive()
)
unloading = (
self._unload_pending
or self.active_model_name != expected_model
or not dispatcher_alive
)
if not unloading:
self._mailboxes[request_id] = mailbox
# When bailing without a mailbox, note whether any OTHER compare request still
# routes through the dispatcher; if none and this call started it, stop it below.
orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes
if unloading:
# A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was
# stopped, then set _unload_pending. The one we just started would otherwise
# linger with no mailboxes, race unload_model's _wait_response for the "unloaded"
# reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop
# it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock:
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
if orphaned_dispatcher:
self._stop_dispatcher()
yield "Error: model is being unloaded"
return
try:
self._send_cmd(cmd)
@ -676,14 +777,18 @@ class InferenceOrchestrator:
return
logger.warning("Timed out draining mailbox after cancel")
def _wait_dispatcher_idle(self) -> None:
def _wait_dispatcher_idle(self) -> bool:
"""Wait for all dispatched requests to complete, then stop dispatcher.
Called by _generate_inner before the _gen_lock path so the dispatcher
thread isn't competing for resp_queue reads.
Returns True if the dispatcher was stopped (all mailboxes drained, or no
dispatcher was running), and False if it was left running because compare
requests were still active after _DISPATCH_IDLE_TIMEOUT.
Called before the _gen_lock path so the dispatcher thread isn't competing
for resp_queue reads.
"""
if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive():
return
return True
# Wait for all mailboxes to be emptied (dispatched requests complete)
deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT
@ -704,8 +809,9 @@ class InferenceOrchestrator:
"leaving dispatcher running for compare requests",
len(self._mailboxes),
)
else:
self._stop_dispatcher()
return False
self._stop_dispatcher()
return True
# ------------------------------------------------------------------
# Public API — same interface as InferenceBackend
@ -772,6 +878,19 @@ class InferenceOrchestrator:
)
for attempt in range(2):
# Stop-loading (/unload -> cancel_load) aborts a load by discarding this
# model's loading marker. cancel_load only kills a live child; if the cancel
# lands before any child exists (GPU placement, or between retries) there is
# nothing to kill, and without this check the loop would spawn a worker and
# load the model after /unload reported it unloaded. Observe removal and stop.
if model_name not in self.loading_models:
logger.info(
"Load for '%s' was cancelled before spawn; not starting a worker",
model_name,
)
self.active_model_name = None
self.models.clear()
return False
logger.info(
"Spawning fresh inference subprocess for '%s' "
"(transformers %s.x, attempt %d/2%s)",
@ -783,6 +902,22 @@ class InferenceOrchestrator:
sub_config["disable_xet"] = disable_xet
self._spawn_subprocess(sub_config)
# A cancel can land after the pre-spawn recheck but while _spawn_subprocess
# is still creating the queues/process. cancel_load runs off the lifecycle
# gate, so its _shutdown_subprocess can see _proc still None and no-op,
# orphaning this fresh worker; the load would then wait for "loaded" and
# publish a model /unload reported unloaded, over a live subprocess nothing
# reaps. Recheck now the child exists and tear it down before publishing.
if model_name not in self.loading_models:
logger.info(
"Load for '%s' was cancelled during spawn; tearing the worker down",
model_name,
)
self._shutdown_subprocess(timeout = 5)
self.active_model_name = None
self.models.clear()
return False
try:
resp = self._wait_response("loaded")
except DownloadStallError:
@ -803,8 +938,31 @@ class InferenceOrchestrator:
)
if resp.get("success"):
# A cancel can land while we were parked in _wait_response above.
# cancel_load (off the lifecycle gate) discards this model's loading
# marker BEFORE its teardown, so a Stop-loading that fired after the
# worker queued "loaded" (which we can still consume during cancel_load's
# shutdown window) shows up here only as the marker's removal. Without
# this recheck we would publish active_model_name/models for a model
# /unload reported cancelled, over a subprocess cancel_load just killed;
# its post-teardown re-clear cannot undo a publish that lands after it
# returns. Observe the removal and abort; cancel_load owns teardown.
if model_name not in self.loading_models:
logger.info(
"Load for '%s' was cancelled while waiting for 'loaded'; "
"not publishing the cancelled model",
model_name,
)
self.active_model_name = None
self.models.clear()
return False
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
# A load always spawns a fresh subprocess holding only this model, so
# mirror that. A lingering stale name would pass unload_model's "not in
# self.models" guard, and the worker's absent-name fallback would unload
# its *active* model, not the already-gone one.
self.models = {}
self.models[self.active_model_name] = {
"is_vision": model_info.get("is_vision", False),
"is_lora": model_info.get("is_lora", False),
@ -837,17 +995,65 @@ class InferenceOrchestrator:
self.models.clear()
raise
def unload_model(self, model_name: str) -> bool:
"""Unload a model from the subprocess."""
if model_name in self.loading_models:
logger.info(
"Cancelling in-flight load for model '%s' by terminating subprocess",
def cancel_load(self, model_name: str) -> bool:
"""Abort an in-flight load by terminating its subprocess.
Returns True if a load for ``model_name`` (matched case-insensitively) was
cancelled, False if nothing was loading under that name. This only tears the
loading subprocess down -- it sends no command to a worker -- so, unlike the
rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle
gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt
a safetensors load that holds the gate for its whole (multi-minute) duration;
a gated cancel could never preempt that load.
"""
target = model_name
if target not in self.loading_models:
target = next(
(m for m in self.loading_models if m.lower() == model_name.lower()),
model_name,
)
self._shutdown_subprocess(timeout = 0.5)
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
if target not in self.loading_models:
return False
logger.info(
"Cancelling in-flight load for model '%s' by terminating subprocess",
target,
)
# Discard the loading marker (and clear local state) BEFORE the teardown, not
# after. cancel_load runs off the lifecycle gate, alongside a load_model that
# rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s
# tearing a live child down and joining the dispatcher), so clearing only after
# leaves a window where load_model reads the marker still set, passes its pre-spawn
# recheck, and loads the model after /unload reported it cancelled. Clear first.
self.loading_models.discard(target)
self.active_model_name = None
self.models.clear()
self._shutdown_subprocess(timeout = 0.5)
# Clear the local mirrors again AFTER the teardown. A racing off-gate load_model
# may still be parked in _wait_response("loaded"): its worker already queued a
# "loaded" reply, so during the shutdown window above (the 0.5s settle before the
# response queue is drained and nulled) that thread can consume it and repopulate
# active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess
# nulls the queue but not the mirrors, so without this second clear /unload reports
# success while the backend still advertises a killed model. The nulled queue lets
# no further "loaded" through, so re-clearing here wipes any repopulation.
self.active_model_name = None
self.models.clear()
return True
def unload_model(self, model_name: str) -> bool:
"""Unload a model from the subprocess."""
# active_model_name can differ in case from the client's raw /unload name (the
# load path canonicalizes casing). Match case-insensitively and use the canonical
# spelling so the guard, unload command, and cleanup below hit the loaded model.
if (
self.active_model_name is not None
and model_name != self.active_model_name
and model_name.lower() == self.active_model_name.lower()
):
model_name = self.active_model_name
# In-flight load: tear its subprocess down (shared loading-cancel logic; no
# worker command sent).
if self.cancel_load(model_name):
return True
if not self._ensure_subprocess_alive():
@ -857,30 +1063,93 @@ class InferenceOrchestrator:
self.active_model_name = None
return True
try:
self._send_cmd(
{
"type": "unload",
"model_name": model_name,
}
)
resp = self._wait_response("unloaded")
# Update local state
# Nothing loaded under this name: don't unload a stale model. The worker falls
# back to unloading its *active* model when the name is absent, so a stale unload
# (lost a race to a concurrent load) would hit the wrong one.
if model_name != self.active_model_name and model_name not in self.models:
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
logger.info("Model '%s' unloaded from subprocess", model_name)
return True
except Exception as exc:
logger.error("Error unloading model '%s': %s", model_name, exc)
# Clear local state anyway
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return False
# The subprocess runs commands sequentially, so a bare unload queues behind a
# running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker
# polls each token), then take _gen_lock as sole resp_queue reader (like GGUF).
#
# Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of
# the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a
# compare request's _start_dispatcher queued behind that stop then observes the
# unload and refuses to spawn a fresh dispatcher that would eat the "unloaded"
# reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet),
# so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock.
with self._dispatcher_lifecycle_lock:
self._unload_pending = True
# Cancelling only the running generation isn't enough: the worker clears
# cancel_event at each generate start, so a queued one would clear it and run the
# outgoing model to completion. drain_event, never cleared, makes any generate
# dequeued during the unload skip.
if self._drain_event is not None:
self._drain_event.set()
try:
self._cancel_generation()
acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT)
if not acquired:
# Wedged worker: tear the subprocess down to free the GPU (next load respawns).
logger.warning(
"Unload: generation did not yield %.1fs after cancel; "
"shutting the inference subprocess down to free the model",
_UNLOAD_GEN_LOCK_TIMEOUT,
)
self._shutdown_subprocess(timeout = 5)
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return True
try:
# Stop the compare-mode dispatcher so it can't consume the "unloaded" reply
# off resp_queue before we do. A dispatched generation bypasses _gen_lock, so
# a wedged one slips past the acquire above; if the dispatcher is still active
# it owns resp_queue and the queued unload hangs _wait_response behind the
# stuck generate. Mirror the wedged locked path: tear the subprocess down.
if not self._wait_dispatcher_idle():
logger.warning(
"Unload: compare-mode dispatcher still active after idle "
"wait; shutting the inference subprocess down to free the model"
)
self._shutdown_subprocess(timeout = 5)
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return True
# Drop stale tokens so they can't be read as the unload reply.
self._drain_queue()
self._send_cmd(
{
"type": "unload",
"model_name": model_name,
}
)
self._wait_response("unloaded")
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
logger.info("Model '%s' unloaded from subprocess", model_name)
return True
except Exception as exc:
logger.error("Error unloading model '%s': %s", model_name, exc)
# Clear local state anyway
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return False
finally:
self._gen_lock.release()
finally:
self._unload_pending = False
if self._drain_event is not None:
self._drain_event.clear()
def generate_chat_response(
self,
@ -899,6 +1168,7 @@ class InferenceOrchestrator:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate response, streaming tokens from subprocess.
@ -908,6 +1178,8 @@ class InferenceOrchestrator:
``stats_holder``: caller-owned dict; on gen_done its "stats" key gets
the worker's usage/timings. Request-scoped to avoid cross-stream reads.
``presence_penalty`` matches the GGUF sampling path (0 disables it).
"""
yield from self._generate_inner(
messages = messages,
@ -926,6 +1198,7 @@ class InferenceOrchestrator:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
stats_holder = stats_holder,
presence_penalty = presence_penalty,
)
def generate_chat_completion_with_tools(
@ -945,6 +1218,7 @@ class InferenceOrchestrator:
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
@ -952,6 +1226,7 @@ class InferenceOrchestrator:
bypass_permissions: bool = False,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
**_unused,
):
"""Run the safetensors agentic tool loop in the parent process,
@ -987,6 +1262,7 @@ class InferenceOrchestrator:
preserve_thinking = preserve_thinking,
# last turn wins, like the GGUF tool loop
stats_holder = stats_holder,
presence_penalty = presence_penalty,
)
if use_adapter is not None:
yield from self.generate_with_adapter_control(
@ -1007,6 +1283,7 @@ class InferenceOrchestrator:
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
nudge_tool_calls = nudge_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
@ -1053,6 +1330,7 @@ class InferenceOrchestrator:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Inner generation logic — sends command to subprocess, yields tokens.
@ -1066,6 +1344,7 @@ class InferenceOrchestrator:
if not self.active_model_name:
yield "Error: No active model"
return
expected_model = self.active_model_name
# Drain any prior compare-mode dispatcher so we can read resp_queue.
self._wait_dispatcher_idle()
@ -1074,6 +1353,14 @@ class InferenceOrchestrator:
# consume and drop each other's token events. Hold _gen_lock across the
# cmd build + send + whole stream so we stay the sole resp_queue reader.
with self._gen_lock:
# Recheck under the lock: an unload we raced may have cleared/swapped the model.
# _unload_pending resets after the lock releases, so it can read False by now;
# the active-model check catches that handoff and a reload that swapped models,
# so we never generate on the wrong one.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield "Error: model is being unloaded"
return
request_id = str(uuid.uuid4())
image_b64 = self._pil_to_base64(image) if image is not None else None
cmd = self._build_generate_cmd(
@ -1087,6 +1374,7 @@ class InferenceOrchestrator:
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
@ -1141,53 +1429,62 @@ class InferenceOrchestrator:
raise RuntimeError("Inference subprocess is not running")
if not self.active_model_name:
raise RuntimeError("No active model")
expected_model = self.active_model_name
request_id = str(uuid.uuid4())
# Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the
# outgoing model once an unload is pending, like the text and audio-input paths.
# Without this a concurrent /audio/generate could run TTS on a model being switched.
with self._gen_lock:
# Recheck under the lock (see _generate_inner): a raced unload/switch may have
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
raise RuntimeError("model is being unloaded")
cmd = {
"type": "generate_audio",
"request_id": request_id,
"text": text,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
request_id = str(uuid.uuid4())
self._send_cmd(cmd)
cmd = {
"type": "generate_audio",
"request_id": request_id,
"text": text,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
# Wait for audio_done or audio_error
deadline = time.monotonic() + 120.0
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout = min(remaining, 1.0))
self._send_cmd(cmd)
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("audio generation"))
continue
deadline = time.monotonic() + 120.0
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout = min(remaining, 1.0))
rtype = resp.get("type", "")
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("audio generation"))
continue
if rtype == "audio_done":
wav_bytes = base64.b64decode(resp["wav_base64"])
sample_rate = resp["sample_rate"]
return wav_bytes, sample_rate
rtype = resp.get("type", "")
if rtype == "audio_error":
raise RuntimeError(resp.get("error", "Audio generation failed"))
if rtype == "audio_done":
wav_bytes = base64.b64decode(resp["wav_base64"])
sample_rate = resp["sample_rate"]
return wav_bytes, sample_rate
if rtype == "error":
raise RuntimeError(resp.get("error", "Unknown error"))
if rtype == "audio_error":
raise RuntimeError(resp.get("error", "Audio generation failed"))
if rtype == "status":
continue
if rtype == "error":
raise RuntimeError(resp.get("error", "Unknown error"))
raise RuntimeError("Timeout waiting for audio generation (120s)")
if rtype == "status":
continue
raise RuntimeError("Timeout waiting for audio generation (120s)")
def generate_whisper_response(
self,
@ -1252,8 +1549,15 @@ class InferenceOrchestrator:
if not self.active_model_name:
yield "Error: No active model"
return
expected_model = self.active_model_name
with self._gen_lock:
# Recheck under the lock (see _generate_inner): a raced unload/switch may have
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield "Error: model is being unloaded"
return
request_id = str(uuid.uuid4())
# numpy array -> list for mp.Queue serialization

80
studio/backend/core/inference/passthrough_healing.py
View file

@ -29,10 +29,25 @@ import os
from collections.abc import Mapping
from typing import Any, Optional
from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal
from core.inference.tool_loop_controller import coerce_tool_arguments
from core.tool_healing import parse_tool_calls_from_text
# Only the formats this healer's parser can promote -- narrower than the loops'
# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare
# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a
# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in.
_HEAL_SIGNALS = (
"<tool_call>",
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
)
def _has_heal_signal(text: str) -> bool:
return any(s in text for s in _HEAL_SIGNALS)
# Read once at import (same convention as the other UNSLOTH_* switches).
_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1"
# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process
@ -44,7 +59,7 @@ def nudge_enabled(request_flag: Optional[bool]) -> bool:
return _NUDGE_DEFAULT if request_flag is None else bool(request_flag)
_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS)
_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS)
# A suspected-but-unclosed tool block larger than this is declared a false
# alarm and flushed, bounding memory on a model rambling XML-lookalike text.
_MAX_HOLD_CHARS = 64 * 1024
@ -198,7 +213,7 @@ def heal_openai_message_events(
if not isinstance(msg, dict) or msg.get("tool_calls"):
return None
content = msg.get("content")
if not isinstance(content, str) or not has_tool_signal(content):
if not isinstance(content, str) or not _has_heal_signal(content):
return None
parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
@ -248,7 +263,7 @@ def heal_openai_message(
def _earliest_signal(buffer: str) -> int:
best = -1
for signal in TOOL_XML_SIGNALS:
for signal in _HEAL_SIGNALS:
index = buffer.find(signal)
if index >= 0 and (best < 0 or index < best):
best = index
@ -275,7 +290,7 @@ def _partial_signal_suffix(buffer: str) -> int:
"""Length of the longest buffer suffix that is a proper prefix of a signal."""
for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1):
tail = buffer[-length:]
if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS):
if any(signal.startswith(tail) for signal in _HEAL_SIGNALS):
return length
return 0
@ -340,9 +355,10 @@ class StreamToolCallHealer:
events.append(("text", emit))
self._buffer = self._buffer[len(self._buffer) - keep :]
return events
# HOLD: handle the FIRST complete block per pass so events keep
# document order (a later declared call must not overtake an
# earlier undeclared one flushing as text).
# HOLD: drain the first contiguous run per pass so events keep document
# order (a later declared call must not overtake an earlier undeclared one
# flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS]
# array of contiguous spans, so later calls in it are not stranded as text.
parsed, spans = parse_tool_calls_from_text(
self._buffer,
id_offset = self._id_offset,
@ -363,26 +379,32 @@ class StreamToolCallHealer:
self._holding = False
continue
return events
start, end = spans[0]
promoted = _promote(
[parsed[0]],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
if start:
events.append(("text", self._buffer[:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
# Drop exactly the promoted markup span; everything else
# (leading text, later blocks) stays and is rescanned.
self._buffer = self._buffer[end:]
else:
# Undeclared or unusable name: its markup is DATA, flush it
# (and anything before it) verbatim, then rescan the rest.
events.append(("text", self._buffer[:end]))
self._buffer = self._buffer[end:]
pos = 0
run_end = spans[0][1]
for order, (call, (start, end)) in enumerate(zip(parsed, spans)):
# Stop at the first gap or incomplete trailing block: leave it for the
# next pass to re-hold and stream incrementally, not flush as text early.
if order and start != run_end:
break
promoted = _promote(
[call],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
# Flush any leading text, then drop the promoted markup span.
if self._buffer[pos:start]:
events.append(("text", self._buffer[pos:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
else:
# Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim.
events.append(("text", self._buffer[pos:end]))
pos = end
run_end = end
# Everything past the drained run (later blocks) stays and is rescanned.
self._buffer = self._buffer[run_end:]
self._holding = False
def finalize(self) -> list:
@ -508,7 +530,7 @@ def nudge_should_retry(
if not message or message.get("tool_calls"):
return False
text = message.get("content")
if not isinstance(text, str) or not has_tool_signal(text):
if not isinstance(text, str) or not _has_heal_signal(text):
return False
return not _heal_would_promote(text, allowed_tools, tools)

49
studio/backend/core/inference/presence_penalty.py Normal file
View file

@ -0,0 +1,49 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Presence-penalty logits helpers for the safetensors/MLX inference paths.
Kept in a dependency-light leaf module (torch + transformers only, no unsloth /
peft) so the pure logic can be imported and unit-tested without pulling in the
full inference backend. ``core.inference.inference`` re-exports these for the
runtime generate paths.
"""
import torch
def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int):
"""OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct
completion token (positions >= prompt_len; prompt excluded, multiplicity
ignored, negatives raise). In place; zero is a no-op."""
if not penalty:
return scores
vocab_size = scores.shape[-1]
for b in range(input_ids.shape[0]):
generated = input_ids[b, prompt_len:]
if generated.numel() == 0:
continue
seen = torch.unique(generated)
# Bound generated ids to the valid range [0, vocab_size). Real completion
# tokens are always in range, so this is a zero-regression safety net that
# drops any stray out-of-range or negative id before indexing (mirrors the
# MLX path's bound). Filtering both ends avoids indexing scores with a
# negative id (which would silently wrap to the wrong row).
seen = seen[(seen >= 0) & (seen < vocab_size)]
if seen.numel():
scores[b, seen] = scores[b, seen] - penalty
return scores
def _make_presence_penalty_processor(penalty: float, prompt_len: int):
"""``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical)."""
if not penalty:
return None
from transformers import LogitsProcessor, LogitsProcessorList
class _PresencePenaltyLogitsProcessor(LogitsProcessor):
@torch.no_grad()
def __call__(self, input_ids, scores):
return apply_presence_penalty(input_ids, scores, penalty, prompt_len)
return LogitsProcessorList([_PresencePenaltyLogitsProcessor()])

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

@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via
``core.inference.tools``.
"""
import bisect
import re
import threading
from typing import Callable, Generator, Optional
@ -21,14 +22,38 @@ from typing import Callable, Generator, Optional
from loggers import get_logger
from core.inference.tool_call_parser import (
_TOOL_ALL_PATS,
_GEMMA_BARE_TC_PREFIX_RE,
_GEMMA_BARE_TC_RE,
_TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS,
_balanced_brace_end,
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
_strip_glm_calls,
_strip_mistral_closed_calls,
_strip_mistral_reasoning,
BUDGET_EXHAUSTED_NUDGE,
MAX_ACT_REPROMPTS,
RAG_MAX_SEARCHES_PER_TURN,
RAG_SEARCH_CAP_NUDGE,
TOOL_XML_SIGNALS,
is_short_intent_without_action,
parse_tool_calls_from_text,
reprompt_to_act_message,
strip_leading_bare_json_call,
strip_llama3_leading_sentinels,
strip_tool_markup,
)
# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated
# pattern lists, so the safetensors streaming strip stays aligned with the parser.
from core.tool_healing import (
_REHEARSAL_TAIL_STRIP_RE,
_strip_bracket_tag_calls,
_think_spans_outside_tool_markup,
apply_tool_strip_patterns,
strip_outside_think,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
coerce_tool_arguments,
@ -50,19 +75,213 @@ logger = get_logger(__name__)
# Buffer cap while disambiguating a possible tool-call prefix.
_MAX_BUFFER_CHARS = 32
# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances.
_MAX_BARE_JSON_BUFFER = 16384
# No grammar constraint here (unlike llama-server's lazy grammar): collapse
# exact-duplicate calls and cap the count so a runaway turn cannot fan out.
_MAX_TOOL_CALLS_PER_TURN = 8
def _active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in active_tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
def _active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in active_tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal;
# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held.
_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?")
def _is_rehearsal_prefix(
stripped: str,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> bool:
"""True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]``
rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space
means prose. Unrestricted mode accepts any identifier; else NAME must be active."""
if not stripped or any(ch.isspace() for ch in stripped):
return False
if unrestricted:
return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None
for name in _active_tool_names(active_tools):
if stripped == name or f"{name}[ARGS]".startswith(stripped):
return True
return False
def _held_rehearsal_tail_len(
text: str,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""Length of a trailing bare tool-name token that may be a split rehearsal call
(``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it
instead of leaking the name. Returns 0 for ordinary prose."""
i = len(text)
while i > 0 and not text[i - 1].isspace():
i -= 1
tail = text[i:]
return (
len(tail)
if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted)
else 0
)
def _rehearsal_name_start(
candidate: str,
signal_pos: int,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding
bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the
signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode)."""
if not candidate.startswith("[ARGS]", signal_pos):
return signal_pos
j = signal_pos
while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"):
j -= 1
if j < signal_pos and (
unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools)
):
return j
return signal_pos
def _earliest_tool_signal(
candidate: str,
signals,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""Index where the turn's first genuine tool-call boundary begins, or -1.
Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal
only when an active tool name (any name in unrestricted mode) precedes it, so a
literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a
real ``NAME[ARGS]`` the boundary is pulled back to NAME."""
best = -1
for sig in signals:
if sig != "[ARGS]":
p = candidate.find(sig)
if p >= 0 and (best < 0 or p < best):
best = p
continue
from_idx = 0
while True:
p = candidate.find("[ARGS]", from_idx)
if p < 0:
break
name_start = _rehearsal_name_start(
candidate, p, active_tools, unrestricted = unrestricted
)
if name_start < p:
# Genuine ``NAME[ARGS]``: the boundary is the start of NAME.
if best < 0 or name_start < best:
best = name_start
break
# Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found.
from_idx = p + len("[ARGS]")
return best
def _has_genuine_tool_signal(
candidate: str,
signals,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> bool:
"""True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``.
Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only
when an active tool name (any in unrestricted mode) precedes it. Mirrors the
``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not
drain inactive-name prose."""
for sig in signals:
if sig == "[ARGS]":
if (
_earliest_tool_signal(
candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted
)
>= 0
):
return True
continue
if sig in candidate:
return True
return False
def strip_tool_markup_streaming(
text: str,
*,
auto_heal_tool_calls: bool = True,
tool_protocol_active: bool = False,
enabled_tool_names: Optional[set] = None,
) -> str:
"""Strip open-ended tool XML from display text without trimming whitespace."""
"""Strip open-ended tool XML from display text without trimming whitespace.
Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so
streaming and final display agree: balanced strips first (nested JSON removed whole),
then the guarded function-XML / GLM scans that close at each call's REAL terminator so
literal markup inside argument values is data and trailing prose survives. Reasoning
``<think>`` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must
not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names``
keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose,
not a call), matching the parse / detection active-tool gate."""
if not (auto_heal_tool_calls or tool_protocol_active):
return text
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
# Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the
# ``<think>`` channel) so raw reasoning does not leak into streamed display; an unclosed
# leading block is held (dropped to EOF) until its closer streams in.
text = _strip_mistral_reasoning(text)
def _seg(segment: str, is_last: bool) -> str:
# Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced
# strips first, then the guarded function-XML / GLM scans, then the regex arms
# (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last
# segment (a bare ``foo[ARGS]`` before <think> is prose). Rehearsal strips are name-gated.
seg = _strip_mistral_closed_calls(segment)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
seg = _strip_function_xml_calls(seg, final = is_last)
seg = _strip_glm_calls(seg, final = is_last)
pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS
for pat in pats:
seg = pat.sub("", seg)
if is_last:
seg = apply_tool_strip_patterns(
seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names
)
return seg
# Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then
# regrows the cumulative text, corrupting append-by-length consumers.
return strip_outside_think(text, _seg)
def _strip_tool_markup_final(
@ -70,10 +289,11 @@ def _strip_tool_markup_final(
*,
auto_heal_tool_calls: bool,
tool_protocol_active: bool = False,
enabled_tool_names: Optional[set] = None,
) -> str:
if not (auto_heal_tool_calls or tool_protocol_active):
return text
return strip_tool_markup(text, final = True)
return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names)
def _status_for_tool(tool_name: str, arguments: dict) -> str:
@ -81,25 +301,76 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str:
return status_for_tool(tool_name, arguments)
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
probe = strip_llama3_leading_sentinels(text.lstrip())
if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
return False
return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
_TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"')
# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional
# render-html card fires for bracket-tag serializations too.
_MISTRAL_RENDER_NAME_RE = re.compile(
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
)
_REHEARSAL_RENDER_NAME_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
def _detect_render_html_tool_start(content: str) -> bool:
"""Return True when the first drained tool call is clearly render_html."""
function_match = _FUNCTION_SIGNAL_RE.search(content)
tool_call_index = content.find("<tool_call>")
if not function_match and tool_call_index < 0:
"""Return True when the FIRST tool call in ``content`` is clearly render_html.
Covers every serialization the loop executes (XML ``<function=>`` / ``<tool_call>``,
Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a
render_html marker inside another call's argument is treated as data. Markers inside
a ``<think>`` / ``[THINK]`` block are dropped since the parser skips them."""
think_spans = _think_spans_outside_tool_markup(content)
_think_starts = [s for s, _e in think_spans]
def _in_think(pos: int) -> bool:
if not think_spans:
return False
i = bisect.bisect_right(_think_starts, pos) - 1
return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1]
def _first_outside(start: int, finder) -> int:
# First occurrence at/after ``start`` that is not inside a think span.
pos = finder(start)
while pos >= 0 and _in_think(pos):
pos = finder(pos + 1)
return pos
candidates: list[tuple[int, str]] = []
for fm in _FUNCTION_SIGNAL_RE.finditer(content):
if not _in_think(fm.start()):
candidates.append((fm.start(), fm.group(1)))
break
tc = _first_outside(0, lambda i: content.find("<tool_call>", i))
if tc >= 0:
nm = _TOOL_CALL_NAME_RE.search(content[tc:])
candidates.append((tc, nm.group(1) if nm else ""))
mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i))
if mt >= 0:
mm = _MISTRAL_RENDER_NAME_RE.match(content, mt)
if mm:
candidates.append((mt, mm.group(1)))
else:
# Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the
# first call through the parser (it reads top-level names).
arr_calls = parse_tool_calls_from_text(content[mt:])
if arr_calls:
candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or ""))
for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content):
if not _in_think(rm.start(1)):
candidates.append((rm.start(1), rm.group(1)))
break
if not candidates:
return False
if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index):
return function_match.group(1) == "render_html"
if tool_call_index >= 0:
name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:])
return bool(name_match and name_match.group(1) == "render_html")
return False
_pos, name = min(candidates, key = lambda c: c[0])
return name == "render_html"
def _coerce_arguments_with_provenance(
@ -149,6 +420,7 @@ def run_safetensors_tool_loop(
execute_tool: Callable[..., str],
cancel_event: Optional[threading.Event] = None,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
@ -188,8 +460,17 @@ def run_safetensors_tool_loop(
for _ev in _auto["events"]:
yield _ev
conversation.extend(_auto["messages"])
# Autoinject ran a KB search outside the controller, so it counts as an
# executed tool for the plan-without-action gate.
rag_autoinjected = bool(_auto)
unrestricted_tools = not tools
# Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the
# ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted.
_enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools))
# Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent
# one-shot), else its repeat is stripped but never drained and the turn ends blank.
_detect_tools = [] if unrestricted_tools else list(tools or [])
tool_controller = ToolLoopController(
tools = None if unrestricted_tools else tools,
auto_heal_tool_calls = auto_heal_tool_calls,
@ -198,6 +479,14 @@ def run_safetensors_tool_loop(
kb_search_count = 0
final_attempt_done = False
next_call_id = 0
reprompt_count = 0
# A denied tool confirmation must not be answered with a plan-without-action
# re-prompt (which would raise the confirmation gate again).
tool_denied = False
# Real tool-call turns completed. Only turns that actually executed a tool count
# against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a
# plan-without-action re-prompt) must not consume budget, matching the GGUF loop.
_executed_tool_iters = 0
def _tool_succeeded(tool_name: str) -> bool:
key_prefix = f"{tool_name}:"
@ -215,9 +504,13 @@ def run_safetensors_tool_loop(
_state_streaming = 1
_state_draining = 2
for iteration in range(max_tool_iterations + 1):
# Reserve re-prompt slots so they don't eat the caller's tool budget.
_extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0
for iteration in range(max_tool_iterations + _extra_iters + 1):
if cancel_event is not None and cancel_event.is_set():
return
# Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
_turn_executed_real_tool = False
if final_attempt_done:
active_tools: list[dict] = []
@ -229,6 +522,8 @@ def run_safetensors_tool_loop(
tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
# Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
_enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools))
detect_state = _state_buffering
content_buffer = ""
@ -304,17 +599,18 @@ def run_safetensors_tool_loop(
if detect_state == _state_streaming:
candidate = cumulative_display + delta
signal_pos = -1
for sig in tool_xml_signals:
p = candidate.find(sig)
if p >= 0 and (signal_pos < 0 or p < signal_pos):
signal_pos = p
# Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is
# pulled back to NAME so the name is not flushed.
signal_pos = _earliest_tool_signal(
candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools
)
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
cleaned_before = strip_tool_markup_streaming(
before_tool,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned_before) > len(last_emitted):
last_emitted = cleaned_before
@ -345,10 +641,20 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives;
# released by later prose or the end-of-stream flush.
if tool_protocol_active:
_hold = _held_rehearsal_tail_len(
cleaned, _detect_tools, unrestricted = unrestricted_tools
)
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
else:
emit = cleaned
if len(emit) > len(last_emitted):
last_emitted = emit
yield {"type": "content", "text": emit}
continue
# BUFFERING: hold until we know it is not a tool call.
@ -366,6 +672,92 @@ def run_safetensors_tool_loop(
if sig.startswith(stripped):
is_prefix = True
break
# Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS]
# counts only with an active NAME so prose is not drained into a no-op.
if sig == "[ARGS]":
if (
_earliest_tool_signal(
stripped,
("[ARGS]",),
_detect_tools,
unrestricted = unrestricted_tools,
)
>= 0
):
is_match = True
break
elif sig.startswith("[") and sig in stripped:
is_match = True
break
# Split rehearsal: hold the bare name until its [ARGS] arrives and matches above.
is_rehearsal_prefix = False
if (
not is_match
and not is_prefix
and tool_protocol_active
and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools)
):
is_prefix = True
is_rehearsal_prefix = True
# Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML
# signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses
# as a call, else stream as content. Non-call text is always recovered downstream.
bare_probe = strip_llama3_leading_sentinels(stripped)
if (
not is_match
and not is_prefix
and tool_protocol_active
and bare_probe.startswith("{")
):
if _balanced_brace_end(bare_probe, 0) is None:
if len(stripped) < _MAX_BARE_JSON_BUFFER:
continue # object still open -- keep buffering
elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names):
# Oversized still-open ENABLED-tool call: stop holding (memory bound) but
# DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams.
detect_state = _state_draining
continue
elif parse_tool_calls_from_text(
content_buffer,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
):
# Closed object that parses as a bare-JSON call -- drain silently.
detect_state = _state_draining
continue
# Closed non-call object (or oversized non-call) -- stream as text.
# Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry:
# buffer it here or it streams raw until the end-of-turn safety net.
# ``(?<!\w)`` keeps "recall:" out; the prefix regex is whitespace-tolerant.
if (
not is_match
and not is_prefix
and tool_protocol_active
and (
"call:".startswith(stripped)
or _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None
or _GEMMA_BARE_TC_RE.match(stripped) is not None
)
):
if _GEMMA_BARE_TC_RE.match(stripped):
detect_state = _state_draining
continue
# A ``call:`` / ``call:partial_name`` prefix with no ``{`` yet: keep
# buffering the variable-length name instead of leaking ``call:longname``.
# Names can exceed 32 chars (OpenAI 64, MCP longer), so a fixed cap would
# flush real calls raw. The prefix regex self-terminates on ordinary prose
# and the ``{`` drains above; bound generously like the bare-JSON path.
if _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None:
if len(stripped) < _MAX_BARE_JSON_BUFFER:
continue
detect_state = _state_draining
continue
if len(stripped) < _MAX_BUFFER_CHARS:
continue # bare "call:" prefix still forming
if is_match:
# Tool signal -- flush any visible prefix before DRAINING
@ -375,6 +767,7 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
@ -398,7 +791,8 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
continue
else:
detect_state = _state_streaming
@ -407,57 +801,121 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Same trailing-name hold as STREAMING for this first flush out of BUFFERING.
if tool_protocol_active:
_hold = _held_rehearsal_tail_len(
cleaned, _detect_tools, unrestricted = unrestricted_tools
)
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
else:
emit = cleaned
if len(emit) > len(last_emitted):
last_emitted = emit
yield {"type": "content", "text": emit}
# Stream finished -- resolve what we collected.
if cancel_event is not None and cancel_event.is_set():
return
if detect_state == _state_buffering:
# Buffer never resolved -- tool XML or plain content?
# Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal
# ``foo[ARGS]{...}`` is not parsed.
stripped = content_buffer.lstrip()
_bare_eos = strip_llama3_leading_sentinels(stripped)
if (
stripped
and tool_protocol_active
and any(sig in stripped for sig in tool_xml_signals)
and _has_genuine_tool_signal(
stripped,
tool_xml_signals,
_detect_tools,
unrestricted = unrestricted_tools,
)
):
detect_state = _state_draining
elif tool_protocol_active and _looks_like_enabled_bare_json(
_bare_eos, _enabled_tool_names
):
# A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary
# JSON answer falls through to the else and streams as content, GGUF parity).
detect_state = _state_draining
else:
# Drain and fall through to STREAMING so the intent re-prompt + safety-net parser
# still fire on short emissions like "Let me search." that never exit BUFFERING.
if content_buffer:
cumulative_display += content_buffer
yield {
"type": "content",
"text": _strip_tool_markup_final(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
yield {"type": "status", "text": ""}
return
cleaned = strip_tool_markup(
cumulative_display, final = True, enabled_tool_names = _enabled_tool_names
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
detect_state = _state_streaming
if detect_state == _state_streaming:
# No tool detected mid-stream -- check for late tool XML.
safety_tc = None
saw_tool_signal = tool_protocol_active and any(
sig in content_accum for sig in tool_xml_signals
# Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's
# strict so plain answers stay untouched. Mirrors GGUF.
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
)
if saw_tool_signal:
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
)
if not safety_tc:
# Final answer: if a literal tool marker in prose was stripped
# during streaming but did not parse as a real call, restore the
# raw cumulative text for core callers. Route-level cleanup can
# still apply the Auto-Heal display policy.
if saw_tool_signal and content_accum:
# Re-prompt once on plan-without-action, before any tool runs
# (GGUF loop parity). The retry is gated on nudge_tool_calls so
# Studio callers (which send True) always nudge, while API callers
# who omit the flag keep today's no-reprompt behavior (opt-in).
stripped_answer = content_accum.strip()
if (
auto_heal_tool_calls
and nudge_tool_calls
and active_tools
and reprompt_count < MAX_ACT_REPROMPTS
and not rag_autoinjected
and not tool_denied
and not any(record.executed for record in tool_controller.history)
and is_short_intent_without_action(stripped_answer)
):
reprompt_count += 1
logger.info(
"Safetensors re-prompt %d/%d: model responded without "
"calling tools (%d chars)",
reprompt_count,
MAX_ACT_REPROMPTS,
len(stripped_answer),
)
conversation.append({"role": "assistant", "content": stripped_answer})
tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
conversation.append(
{
"role": "user",
"content": reprompt_to_act_message(tool_hint),
}
)
# Empty status clears the badge and resets the route's
# per-turn text cursor before the re-prompted turn streams.
yield {"type": "status", "text": ""}
continue
# Final answer. If a literal tool marker in prose was buffered but
# never parsed as a call, restore the raw text so the prose surfaces
# in full; route-level cleanup still applies the Auto-Heal policy.
if content_accum and any(sig in content_accum for sig in tool_xml_signals):
yield {"type": "content", "text": content_accum}
else:
# Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real
# prose, release it.
final_clean = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(final_clean) > len(last_emitted):
yield {"type": "content", "text": final_clean}
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
@ -465,31 +923,41 @@ def run_safetensors_tool_loop(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
enabled_tool_names = _enabled_names_gate,
)
logger.info(
"Safetensors safety net: parsed %d tool call(s) from streamed content",
len(tool_calls),
)
else:
# DRAINING: parse tool calls out of full content.
# DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the
# ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to
# drain here: a spent one-shot (render_html) is off the active list but its re-emitted
# ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of
# being dropped into a blank continuation.
tool_calls = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_names_gate,
)
if not tool_calls:
# Parser found nothing. Auto-Heal-enabled display cleanup
# strips unparseable tool XML; disabled Auto-Heal preserves
# the raw text so literal/malformed markup stays visible.
if content_accum:
yield {
"type": "content",
"text": _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
_drain_text = _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
enabled_tool_names = _enabled_tool_names,
)
# Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment
# (plain JSON answers are left untouched); off keeps it visible per the strict contract.
if tool_protocol_active and auto_heal_tool_calls:
_drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names)
if _drain_text:
yield {"type": "content", "text": _drain_text}
if provisional_render_html_started and not provisional_resolved:
provisional_resolved = True
yield {
@ -505,10 +973,14 @@ def run_safetensors_tool_loop(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
enabled_tool_names = _enabled_names_gate,
)
if tool_calls:
next_call_id += len(tool_calls)
# Strip a leading bare-JSON call from the kept content so it isn't replayed as text or
# next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers.
content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names)
if final_attempt_done:
# Final-answer turn re-called a tool -- stop the loop.
@ -517,6 +989,27 @@ def run_safetensors_tool_loop(
yield {"type": "status", "text": ""}
return
# Collapse exact-duplicate calls and cap the count (runaway-turn guard).
if tool_calls:
seen_keys: set = set()
deduped: list = []
for _tc in tool_calls:
_fn = _tc.get("function", {}) or {}
_key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
if _key in seen_keys:
continue
seen_keys.add(_key)
deduped.append(_tc)
if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN:
break
if len(deduped) != len(tool_calls):
logger.info(
"Safetensors: collapsed %d repeated tool call(s) in one turn to %d",
len(tool_calls),
len(deduped),
)
tool_calls = deduped
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
@ -593,6 +1086,7 @@ def run_safetensors_tool_loop(
"result": TOOL_REJECTED_MESSAGE,
"provenance": decision.provenance,
}
tool_denied = True
denied_message = {
"role": "tool",
"name": decision.tool_name,
@ -634,6 +1128,8 @@ def run_safetensors_tool_loop(
completion = tool_controller.record_result(decision, result)
if provisional_match:
provisional_resolved = True
# A tool ran this turn, so it counts against the caller's budget.
_turn_executed_real_tool = True
yield completion.tool_end_event()
conversation.append(completion.tool_message())
@ -646,7 +1142,11 @@ def run_safetensors_tool_loop(
if not unrestricted_tools and not tool_controller.active_tools():
final_attempt_done = True
continue
if iteration + 1 >= max_tool_iterations and not final_attempt_done:
# Count only turns that executed a tool against the cap; a no-op correction turn doesn't
# consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity).
if _turn_executed_real_tool:
_executed_tool_iters += 1
if _executed_tool_iters >= max_tool_iterations and not final_attempt_done:
# Budget exhausted; nudge a final plain answer.
final_attempt_done = True
conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE})

2745
studio/backend/core/inference/tool_call_parser.py
View file

File diff suppressed because it is too large Load diff

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

@ -406,6 +406,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
)
def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool:
"""Skip a generate queued behind a cancelled one during an unload.
The parent sets ``drain_event`` for the whole unload. Because the parent's
per-token ``cancel_event`` is cleared at the start of every generate, a cancel
set while this generate was still queued would otherwise be lost when it is
dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so
the parent's stream/mailbox drains fast and the switch stays fast, and report
the generate was skipped so the caller does not clear the cancel or run it.
"""
if drain_event is None or not drain_event.is_set():
return False
request_id = cmd.get("request_id", "")
logger.info("Skipping generate for request %s: unload draining", request_id)
_send_response(
resp_queue,
{
"type": "gen_done",
"request_id": request_id,
"cancelled": True,
"stats": None,
},
)
return True
def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"""Handle a generate command: stream tokens back via resp_queue.
@ -431,6 +457,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"min_p": cmd.get("min_p", 0.0),
"max_new_tokens": cmd.get("max_new_tokens", 256),
"repetition_penalty": cmd.get("repetition_penalty", 1.0),
"presence_penalty": cmd.get("presence_penalty", 0.0),
"cancel_event": cancel_event,
}
@ -632,7 +659,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
)
def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None:
def run_inference_process(
*,
cmd_queue: Any,
resp_queue: Any,
cancel_event,
config: dict,
drain_event = None,
) -> None:
"""Subprocess entrypoint. Persistent — runs the command loop until shutdown.
Args:
@ -640,6 +674,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
resp_queue: mp.Queue for sending responses to parent.
cancel_event: mp.Event the parent sets to cancel generation.
config: Initial configuration dict with model info.
drain_event: mp.Event the parent sets for the duration of an unload. Unlike
cancel_event (cleared at the start of every generate), it is never cleared
here, so a generate still queued behind a cancelled one is skipped rather
than run — the cancel survives the queue handoff.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
@ -715,7 +753,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
cmd_type = cmd.get("type", "")
try:
if cmd_type == "generate":
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
cancel_event.clear()
# Re-check the drain after clearing: the parent sets drain_event
# then cancel_event for an unload, so if that pair landed between
# the check above and this clear, the clear just erased the unload's
# cancel. Skip here so the outgoing model is not run to completion,
# which would stall the switch until the dispatcher idle-timeout.
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "load":
if backend.active_model_name:
@ -918,7 +965,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
try:
if cmd_type == "generate":
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
cancel_event.clear()
# Re-check the drain after clearing: the parent sets drain_event then
# cancel_event for an unload, so if that pair landed between the check
# above and this clear, the clear just erased the unload's cancel. Skip
# here so the outgoing model is not run to completion, which would stall
# the switch until the dispatcher idle-timeout tears the subprocess down.
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "load":

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

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

721
studio/backend/core/tool_healing.py
View file

@ -1,38 +1,137 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge
# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026
# Antoine Zambelli, used under the MIT License.
"""Lightweight tool-call XML parsing and stripping helpers.
"""Lightweight tool-call parsing and stripping helpers.
External inference servers import this module without pulling in the inference
orchestrator, structlog, httpx, or the rest of the studio backend.
orchestrator, structlog, httpx, or the rest of the studio backend. Kept in
lockstep with ``core/inference/tool_call_parser.py`` so those servers
(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any
change here must also land there.
Handles these serializations (see ``parse_tool_calls_from_text``):
* ``<tool_call>{json}</tool_call>``
* ``<|tool_call>call:name{...}<tool_call|>`` (Gemma)
* ``<function=name><parameter=k>v</parameter></function>``
* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback)
* ``name[ARGS]{json}`` (reasoning-model rehearsal)
"""
# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9).
from __future__ import annotations
import bisect
import json
import re
# Pre-compiled patterns for tool XML stripping. The hyphen in the name
# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues,
# issue-number) parse alongside the built-ins.
# One nesting level in the strip regexes; deeper may leak markup (still parsed).
_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed =
# complete body, tail = truncated; ``(?<!\[CALL_ID\])`` keeps the v11 call-id from reading as a name.
_REHEARSAL_CLOSED_STRIP_RE = re.compile(
r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*" + _BRACKETED_JSON_ONE_LEVEL, re.DOTALL
)
_REHEARSAL_TAIL_STRIP_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?:\{.*)?$", re.DOTALL)
# Tool-XML strip patterns; hyphen in the name class covers dashed MCP names.
# Closed-pair patterns are named so _PAT_REQUIRED_TOKEN can skip a doomed lazy rescan when
# the close token is absent: an unguarded ``<tag>.*?</tag>`` rescans to EOF from every opener
# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass.
_TC_JSON_CLOSED_PAT = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL)
_TC_FUNC_CLOSED_PAT = re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL)
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL),
_TC_JSON_CLOSED_PAT,
_TC_GEMMA_CLOSED_PAT,
re.compile(r"<tool_call\|>"),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
_TC_FUNC_CLOSED_PAT,
# Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata.
re.compile(
r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*"
+ _BRACKETED_JSON_ONE_LEVEL,
re.DOTALL,
),
_REHEARSAL_CLOSED_STRIP_RE,
# Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind.
re.compile(r"\[/TOOL_CALLS\]"),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF
# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by
# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data).
_TOOL_OPEN_XML_TAIL_PATS = [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<\|tool_call>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
_TOOL_ALL_PATS = (
_TOOL_CLOSED_PATS
+ _TOOL_OPEN_XML_TAIL_PATS
+ [
re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL),
_REHEARSAL_TAIL_STRIP_RE,
]
)
# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None.
_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE})
# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument
# data cannot make the helper truncate the block and its tail.
_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT]
# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every
# opener; skip that doomed (quadratic) pass. Shared by both strip helpers.
_PAT_REQUIRED_TOKEN = {
_TC_JSON_CLOSED_PAT: "</tool_call>",
_TC_GEMMA_CLOSED_PAT: "<tool_call|>",
_TC_FUNC_CLOSED_PAT: "</function>",
}
def strip_tool_patterns(text: str, patterns) -> str:
"""Apply ``patterns`` in order, skipping closed-pair passes with no close token."""
for pat in patterns:
token = _PAT_REQUIRED_TOKEN.get(pat)
if token is not None and token not in text:
continue
text = pat.sub("", text)
return text
def apply_tool_strip_patterns(
text: str,
patterns,
enabled_tool_names = None,
) -> str:
"""Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern
strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is
``None``); every other pattern is removed unconditionally. A closed-pair pattern whose
close token is absent is skipped so an unclosed-marker stream stays linear."""
for pat in patterns:
token = _PAT_REQUIRED_TOKEN.get(pat)
if token is not None and token not in text:
continue
if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS:
text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text)
else:
text = pat.sub("", text)
return text
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{")
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it.
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>[^\S\n]*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_GEMMA_QUOTE = '<|"|>'
_PARAM_CLOSE_TAG = "</parameter>"
@ -43,7 +142,62 @@ _FUNC_CLOSE_TAG = "</function>"
# must be identifier-shaped (start with a letter or underscore); a comma
# followed by digits-then-colon is value text such as a timestamp or ratio
# (`meet at 10:00, 11:00 tomorrow`), not a new key.
_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:")
_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:")
# A candidate starting inside a think block is a rehearsal (block kept so literal tags in
# real args survive); ``$`` accepts an unclosed block mid-stream.
_THINK_TAG_RE = re.compile(r"<think>.*?(?:</think>|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL)
# Bare open/close markers for prefilled-reasoning turns (template opens <think> in the prompt).
_THINK_OPEN_RE = re.compile(r"<think>|\[THINK\]")
_THINK_CLOSE_RE = re.compile(r"</think>|\[/THINK\]")
# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects.
_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)")
# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name,
# and hyphens keep dashed MCP names whole.
_MISTRAL_BRACKET_RE = re.compile(
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
)
# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id
# from being taken as the function name.
_REHEARSAL_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
# Above this size skip the balanced scan; the linear regex catch-all bounds pathological output.
_MAX_BRACKET_SCAN_CHARS = 1_000_000
def _balanced_json_span(text: str, start: int) -> int | None:
"""Return the end index of a balanced JSON object opening at ``start``,
or ``None`` if the braces don't balance. Honors escapes and strings.
"""
if start >= len(text) or text[start] != "{":
return None
depth = 0
in_string = False
escape = False
for j in range(start, len(text)):
ch = text[j]
if escape:
escape = False
continue
if ch == "\\":
escape = True
continue
if in_string:
if ch == '"':
in_string = False
continue
if ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return j
return None
def _balanced_brace_end(
@ -109,6 +263,94 @@ def _balanced_bracket_end(src: str, start: int) -> int:
return -1
def _decode_array_items(text: str, body_start: int, body_end: int):
"""Return ``(objs, ends)`` for each top-level element of the JSON array between
``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded
object and its absolute exclusive end offset.
Decoding element-by-element with ``raw_decode`` tolerates the comma-less object
separators the repo's own Mistral/Ollama multi-call templates emit
(``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the
whole body rejects that form and would drop every call. The ends also tile the
region across the calls' spans so a with_spans consumer strips each exactly once."""
decoder = json.JSONDecoder()
objs: list = []
ends: list[int] = []
i = text.find("[", body_start)
if i < 0:
return objs, ends
i += 1
while i < body_end:
while i < body_end and text[i] in " \t\r\n,":
i += 1
if i >= body_end or text[i] == "]":
break
try:
obj, rel = decoder.raw_decode(text[i:body_end])
except (json.JSONDecodeError, ValueError):
break
i += rel
objs.append(obj)
ends.append(i)
return objs, ends
def _iter_bracket_spans(
text: str,
start: int = 0,
enabled_tool_names = None,
):
"""Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag
call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is
``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11
[CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}).
``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous
bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a
prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit
[TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric.
Balance-only (no JSON validation) so strip and parse share one scan. The cursor
jumps past each consumed span, so a marker inside consumed JSON is never
re-matched and each regex re-searches only once its match falls behind: linear."""
n = len(text)
specs = (
("array", _MISTRAL_ARRAY_RE),
("name", _MISTRAL_BRACKET_RE),
("rehearsal", _REHEARSAL_RE),
)
nexts = {kind: rx.search(text, start) for kind, rx in specs}
cursor = start
while cursor < n:
for kind, rx in specs:
m = nexts[kind]
if m is not None and m.start() < cursor:
nexts[kind] = rx.search(text, cursor)
live = [(kind, m) for kind, m in nexts.items() if m is not None]
if not live:
return
kind, m = min(live, key = lambda km: km[1].start())
if kind == "array":
end = _balanced_bracket_end(text, m.end())
end = None if end < 0 else end
else:
end = _balanced_json_span(text, m.end())
if end is None:
# Truncated body: skip and keep scanning; the caller's catch-all strips the tail.
cursor = m.end()
continue
if (
kind == "rehearsal"
and enabled_tool_names is not None
and m.group(1) not in enabled_tool_names
):
# Inactive-name rehearsal is prose: advance past its body without yielding.
cursor = end + 1
continue
yield (m.start(), end + 1, kind, m)
cursor = end + 1
def _split_top_level_commas(src: str) -> list:
"""Split on commas that are not inside a nested ``[]``/``{}`` or a string."""
parts: list[str] = []
@ -223,7 +465,7 @@ def _quote_gemma_object_keys(src: str) -> str:
while i < len(src) and src[i].isspace():
i += 1
key_name_start = i
while i < len(src) and (src[i].isalnum() or src[i] in "_-"):
while i < len(src) and (src[i].isalnum() or src[i] in "_-."):
i += 1
key_name = src[key_name_start:i]
colon_pos = i
@ -267,7 +509,8 @@ def _quote_gemma_object_keys(src: str) -> str:
json.loads(raw.strip())
parts.append(raw)
except (json.JSONDecodeError, ValueError):
parts.append(json.dumps(raw.strip()) if raw.strip() else raw)
# Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}.
parts.append(json.dumps(raw.strip()))
else:
parts.append(src[key_start:i])
return "".join(parts)
@ -291,9 +534,99 @@ def _inside_open_parameter(content: str, pos: int) -> bool:
last_param_start = match.start()
if last_param_start < 0:
return False
last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos)
last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos)
return last_param_start > max(last_param_close, last_func_close)
# The parameter's OWN close tag decides: if it closes after ``pos`` the position is
# argument data (even across literal function closes); an unclosed one falls back to func close.
own_close = content.find(_PARAM_CLOSE_TAG, last_param_start)
if own_close >= 0:
return own_close > pos
func_close = content.find(_FUNC_CLOSE_TAG, last_param_start)
return func_close < 0 or pos < func_close
def _func_close_index(content: str, body_start: int, body: str) -> int:
"""Index in ``body`` of the first ``</function>`` that is not argument
data (not inside an open parameter value); -1 when every close is data.
Taking the LAST close swallowed prose between the real close and a
literal ``</function>`` mentioned later in the answer."""
idx = body.find(_FUNC_CLOSE_TAG)
while idx >= 0:
if not _inside_open_parameter(content, body_start + idx):
return idx
idx = body.find(_FUNC_CLOSE_TAG, idx + 1)
return -1
def _trim_param_value(val: str) -> str:
"""Trim the single wrapping newline the chat template adds around an XML
parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed
code/diff argument indentation)."""
if val.startswith("\n"):
val = val[1:]
if val.endswith("\n"):
val = val[:-1]
return val
def _marker_coverage(content: str, markers) -> list[tuple[int, int]]:
"""Coverage ``[start, end]`` per marker, used to skip markers that are another
call's data. Closes pair to markers via a per-format stack so an inner close
is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with
a paired close cover through it (markers before the close are data); balanced
without one cover only the braces, so a later sibling is still recovered."""
n = len(content)
brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0]
events = [] # (position, order) with order 0 = braces-done, 1 = close marker
for idx, (_start, brace_end, _kind, _m) in enumerate(markers):
if brace_end >= 0:
events.append((brace_end, 0, _kind, idx))
for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)):
for cm in close_re.finditer(content):
# A close inside another call's balanced braces is quoted data; it
# must not pop an earlier close-less marker and swallow a sibling.
if any(s < cm.start() < be for s, be in brace_regions):
continue
events.append((cm.start(), 1, kind, cm.end()))
events.sort(key = lambda e: (e[0], e[1]))
waiting = {"json": [], "gemma": []}
close_end_for: dict[int, int] = {}
for _pos, order, kind, payload in events:
if order == 0:
waiting[kind].append(payload) # marker index, now awaiting its close
elif waiting[kind]:
close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here
coverage = []
for idx, (start, brace_end, _kind, _m) in enumerate(markers):
if brace_end < 0:
coverage.append((start, n))
elif idx in close_end_for:
coverage.append((start, close_end_for[idx]))
else:
coverage.append((start, brace_end))
return coverage
def _build_markers(content: str):
"""JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document
order; ``brace_end < 0`` marks an unbalanced (to-EOF) open."""
markers = []
for start_re, gemma, kind in (
(_TC_JSON_START_RE, False, "json"),
(_TC_GEMMA_START_RE, True, "gemma"),
):
for m in start_re.finditer(content):
if _inside_open_parameter(content, m.start()):
continue
brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma)
markers.append((m.start(), brace_end, kind, m))
markers.sort(key = lambda c: c[0])
return markers
def marker_coverage(content: str) -> list[tuple[int, int]]:
"""Coverage spans of JSON/Gemma tool markers so other parsers can treat markup
inside a marker's coverage (even a marker that failed to parse) as that call's
data rather than a sibling call."""
return _marker_coverage(content, _build_markers(content))
def parse_tool_calls_from_text(
@ -301,6 +634,7 @@ def parse_tool_calls_from_text(
*,
id_offset: int = 0,
allow_incomplete: bool = True,
enabled_tool_names = None,
with_spans: bool = False,
):
"""Parse OpenAI-format tool calls from model text.
@ -309,55 +643,61 @@ def parse_tool_calls_from_text(
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<|tool_call>call:web_search{query:"..."}<tool_call|>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
[TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback)
web_search[ARGS]{"query":"..."} (reasoning-model rehearsal)
A call rehearsed inside a ``<think>`` / ``[THINK]`` block is skipped, not
executed; the block is kept so a literal tag in a real argument is preserved.
With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]``
is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup
in ``content`` (including its close tag when present), so a caller can
remove exactly the parsed markup and keep every other byte intact.
"""
# Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a
# think marker opening inside a call is argument data (excluded from spans).
_think_spans = _think_spans_outside_tool_markup(content)
_think_starts = [s for s, _e in _think_spans]
def _in_think(pos: int) -> bool:
# Spans are ordered and non-overlapping; bisect gives O(log M) per candidate.
i = bisect.bisect_right(_think_starts, pos) - 1
return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1]
tool_calls: list[dict] = []
call_spans: list[tuple] = []
# Collect every supported call format with spans, then emit in document
# order. A marker inside another call's argument string is data, not a
# separate executable call.
parsed_items = [] # (start, span_end, name, arguments)
candidates = [] # (start, brace_end, kind, match)
for m in _TC_JSON_START_RE.finditer(content):
if _inside_open_parameter(content, m.start()):
continue
end = _balanced_brace_end(content, m.end() - 1)
if end >= 0:
candidates.append((m.start(), end, "json", m))
for m in _TC_GEMMA_START_RE.finditer(content):
if _inside_open_parameter(content, m.start()):
continue
end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True)
if end >= 0:
candidates.append((m.start(), end, "gemma", m))
candidates.sort(key = lambda c: c[0])
candidate_spans = [(s, e) for s, e, _kind, _m in candidates]
for idx, (start, end, kind, m) in enumerate(candidates):
if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx):
# Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside
# another call's coverage (even one that failed to parse) is data, not executed. A
# marker opening inside a think block is a rehearsal and is skipped.
parsed_items = [] # (start, span_end, name, arguments) in document order
markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])]
coverage = _marker_coverage(content, markers)
for idx, (start, brace_end, kind, m) in enumerate(markers):
if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx):
continue
if brace_end < 0:
continue # unclosed: not parseable; the fallback still excludes its XML
if not allow_incomplete:
tail = content[end + 1 :].lstrip()
tail = content[brace_end + 1 :].lstrip()
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
if close_re.match(tail) is None:
continue
try:
if kind == "json":
obj = json.loads(content[m.end() - 1 : end + 1])
obj = json.loads(content[m.end() - 1 : brace_end + 1])
name = obj.get("name", "")
arguments = obj.get("arguments", {})
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes).
arguments = obj.get("arguments")
if arguments is None:
arguments = obj.get("parameters", {})
if isinstance(arguments, dict):
arguments = json.dumps(arguments)
else:
name = m.group(1)
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end]))
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))
except (json.JSONDecodeError, ValueError):
continue
span_end = end + 1
span_end = brace_end + 1
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
ws = len(content[span_end:]) - len(content[span_end:].lstrip())
close_m = close_re.match(content, span_end + ws)
@ -369,7 +709,8 @@ def parse_tool_calls_from_text(
fm
for fm in _TC_FUNC_START_RE.finditer(content)
if not _inside_open_parameter(content, fm.start())
and not any(s <= fm.start() <= e for s, e in candidate_spans)
and not _in_think(fm.start())
and not any(s <= fm.start() < e for s, e in coverage)
]
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
@ -382,7 +723,7 @@ def parse_tool_calls_from_text(
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
close_idx = body.rfind(_FUNC_CLOSE_TAG)
close_idx = _func_close_index(content, body_start, body)
if close_idx >= 0:
span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
body = body[:close_idx]
@ -404,7 +745,7 @@ def parse_tool_calls_from_text(
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
arguments[pm.group(1)] = _trim_param_value(val)
else:
valid_params = True
for pidx, pm in enumerate(param_starts):
@ -422,7 +763,7 @@ def parse_tool_calls_from_text(
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
arguments[param_name] = _trim_param_value(val)
if not valid_params:
continue
@ -444,19 +785,293 @@ def parse_tool_calls_from_text(
}
)
call_spans.append((start, span_end))
# Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document
# order, so a Mistral call and a rehearsal in one message both parse.
if not tool_calls:
for start, end, kind, m in _iter_bracket_spans(
content, enabled_tool_names = enabled_tool_names
):
if _in_think(start):
continue
# Extend the region over an immediately-following v11 closer so with_spans consumers strip it too.
closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:])
region_end = end + closer.end() if closer else end
if kind == "array":
# Decode elements individually (comma-tolerant): one json.loads of the whole
# body rejects the comma-less multi-call arrays Mistral/Ollama templates emit.
payload, item_ends = _decode_array_items(content, m.end(), end)
if not payload:
continue
# Tile the region so every byte belongs to exactly one span; a with_spans consumer
# keeps skipped bytes visible and strips promoted markup exactly once.
tile_start = start
last_span_idx = -1
for item_idx, item in enumerate(payload):
if not isinstance(item, dict) or "name" not in item:
continue
args = item.get("arguments", {})
if isinstance(args, str):
# ``arguments`` may itself be a JSON string (OpenAI spec).
try:
args = json.loads(args)
except (json.JSONDecodeError, ValueError):
pass
if not isinstance(args, (dict, str)):
# ``"arguments": null`` (or any non-object scalar) becomes {} like the
# <tool_call> path, not the string "null" auto-heal would mangle to
# a bogus {"query":"null"}.
args = {}
tool_calls.append(
{
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": item.get("name", ""),
# A bare scalar string stays raw (like the <tool_call> path);
# json.dumps would double-encode it so the arg healer wraps
# "weather" with its literal quotes.
"arguments": args if isinstance(args, str) else json.dumps(args),
},
}
)
item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end
last_span_idx = len(call_spans)
call_spans.append((tile_start, item_end))
tile_start = item_end
if last_span_idx >= 0:
tile_start, _tile_end = call_spans[last_span_idx]
call_spans[last_span_idx] = (tile_start, region_end)
else:
try:
payload = json.loads(content[m.end() : end])
except (json.JSONDecodeError, ValueError):
continue
if not isinstance(payload, dict):
continue
tool_calls.append(
{
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": m.group(1),
"arguments": json.dumps(payload),
},
}
)
call_spans.append((start, region_end))
if with_spans:
return tool_calls, call_spans
return tool_calls
def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str:
"""Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one
balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex
left two-level args behind). Truncated tails go to the caller's catch-all. Linear.
``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None
strips every span)."""
if len(text) > _MAX_BRACKET_SCAN_CHARS:
return text
out: list[str] = []
cursor = 0
for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names):
out.append(text[cursor:start])
cursor = end
out.append(text[cursor:])
return "".join(out)
def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]:
"""Spans of tool-call markup, so a literal <think>/[THINK] inside a call's args is
stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket
calls and an unclosed XML call (run via allow_incomplete); without the open-ended
span the unclosed call's markup would leak after execution."""
# Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan
# to EOF from every opener (quadratic on a stream of unclosed openers).
spans = [
m.span()
for pat in _TOOL_CLOSED_PATS
if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text)
for m in pat.finditer(text)
]
spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text))
# An unclosed opener is a real incomplete call only outside closed/bracket spans.
for pat in _TOOL_OPEN_XML_TAIL_PATS:
for m in pat.finditer(text):
if not any(s <= m.start() < e for s, e in spans):
spans.append(m.span())
return spans
def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]:
"""<think>/[THINK] block spans, minus any whose opening marker sits INSIDE a
tool-call span (that tag is argument data, not reasoning). Keeping it would drop a
real call after it as rehearsed and leak the call's markup. START tested only, so
a greedy unclosed <think> past the call is still that call's argument data."""
think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)]
call_spans = _tool_call_markup_spans(text)
# Prefilled reasoning: the template opens <think> in the prompt, so add a leading span
# (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe.
close = _THINK_CLOSE_RE.search(text)
if close is not None:
opener = _THINK_OPEN_RE.search(text)
if (
(opener is None or close.start() < opener.start())
and not any(cs <= close.start() < ce for cs, ce in call_spans)
and any(cs >= close.end() for cs, ce in call_spans)
):
think_spans = [(0, close.end())] + think_spans
if not think_spans:
return think_spans
if not call_spans:
return think_spans
return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)]
def strip_outside_think(text: str, strip_segment) -> str:
"""Apply ``strip_segment(segment, is_last)`` to visible text around <think>/[THINK]
blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal).
``is_last`` is True only after the final block, so trailing-tail patterns apply
only there. Shared by every strip path so they stay consistent."""
# A think marker opening inside a complete call is argument text; excluding it lets the
# stripper see the whole call. START-tested, so an unclosed match stays argument data.
think_spans = _think_spans_outside_tool_markup(text)
if not think_spans:
return strip_segment(text, True)
pieces: list[str] = []
prev = 0
for s, e in think_spans:
pieces.append(strip_segment(text[prev:s], False))
pieces.append(text[s:e])
prev = e
pieces.append(strip_segment(text[prev:], True))
return "".join(pieces)
def _strip_gemma_native_spans(text: str, *, final: bool) -> str:
"""Remove complete Gemma-native spans, brace/quote-balanced so a literal
``<tool_call|>`` in a quoted argument cannot truncate the span. An incomplete
span is dropped to EOF when ``final``, else kept (still streaming)."""
out: list[str] = []
cursor = 0
for match in _TC_GEMMA_START_RE.finditer(text):
start = match.start()
if start < cursor:
continue
brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
if brace_end < 0:
# Unbalanced: nothing completes from here on. Drop the rest if final,
# else keep it; stop either way (rescanning would be quadratic).
if final:
out.append(text[cursor:start])
cursor = len(text)
break
# Junk between } and <tool_call|> is malformed-call markup: strip through
# the close, keep text after it. No close anywhere means stop (linear).
close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
if close is None:
if final:
out.append(text[cursor:start])
cursor = len(text)
break
out.append(text[cursor:start])
cursor = close.end()
out.append(text[cursor:])
return "".join(out)
def _gemma_span_ranges(text: str) -> list:
"""``(start, end)`` of each complete Gemma-native span; same walk as
``_strip_gemma_native_spans`` without stripping."""
ranges: list[tuple] = []
cursor = 0
for match in _TC_GEMMA_START_RE.finditer(text):
start = match.start()
if start < cursor:
continue
brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
if brace_end < 0:
break
close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
if close is None:
break
ranges.append((start, close.end()))
cursor = close.end()
return ranges
def _strip_closed_blocks_outside_gemma(text: str) -> str:
"""Closed JSON/function pre-pass that skips matches starting inside a complete
Gemma span: deleting across the span boundary would mangle the Gemma close and
truncate the tail. A skipped match resumes at the covering span's end, so a
real function-XML call after the span is still stripped."""
ranges = _gemma_span_ranges(text)
if not ranges:
return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS)
for pat in _TOOL_CLOSED_BLOCK_PATS:
token = _PAT_REQUIRED_TOKEN.get(pat)
if token is not None and token not in text:
continue
out: list[str] = []
pos = 0
while True:
m = pat.search(text, pos)
if m is None:
out.append(text[pos:])
break
covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None)
if covering is not None:
out.append(text[pos : covering[1]])
pos = covering[1]
continue
out.append(text[pos : m.start()])
pos = m.end()
new_text = "".join(out)
if new_text != text:
text = new_text
ranges = _gemma_span_ranges(text)
return text
def _strip_markup_segment(
text: str,
*,
final: bool,
enabled_tool_names = None,
) -> str:
# Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth,
# rehearsal name-gated); then the quote-aware Gemma-native passes so a literal
# <tool_call|> in an argument cannot truncate a block; finally the regex XML/tail sweeps.
text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names)
text = _strip_closed_blocks_outside_gemma(text)
text = _strip_gemma_native_spans(text, final = final)
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names)
def strip_tool_call_markup(
text: str,
*,
final: bool = False,
enabled_tool_names = None,
) -> str:
"""Strip tool-call XML markup from text.
When ``final`` is False, only fully closed tool-call blocks are removed.
When ``final`` is True, trailing incomplete tool-call blocks are removed
too, and the result is stripped of surrounding whitespace.
``<think>`` / ``[THINK]`` reasoning is preserved verbatim (see
``strip_outside_think``); the trailing-tail patterns apply only after the
last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}``
example visible (it is prose, not a call) so display cleanup matches detection.
"""
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
for pat in patterns:
text = pat.sub("", text)
return text.strip() if final else text
result = strip_outside_think(
text,
lambda seg, is_last: _strip_markup_segment(
seg, final = final and is_last, enabled_tool_names = enabled_tool_names
),
)
return result.strip() if final else result

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

@ -177,6 +177,7 @@ class GenerateRequest(BaseModel):
temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature")
top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling")
top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling")
min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling")
max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate")
repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty")
presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty")
@ -1146,7 +1147,8 @@ class CompletionMessage(BaseModel):
"""The assistant's complete response message."""
role: Literal["assistant"] = "assistant"
content: str
# ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise.
content: Optional[str] = None
refusal: Optional[str] = None
reasoning_content: Optional[str] = None
tool_calls: Optional[list[dict]] = None

6
studio/backend/requirements/single-env/overrides-darwin-arm64.txt
View file

@ -10,3 +10,9 @@ transformers>=4.57.6
# anyio that also ImportErrors on TaskHandle and 500s the server. An override
# wins the fight, so force one consistent <4.14 here too.
anyio<4.14.0
# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights
# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242.
# The override also drops it from transitive resolution; keep the >=0.22.0 floor
# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it.
mlx-lm>=0.22.0,!=0.31.3

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

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

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

File diff suppressed because it is too large Load diff

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

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

79
studio/backend/run.py
View file

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

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

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

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

@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO
from types import SimpleNamespace
def _emitter_client_text(events: list[str]) -> str:
"""Concatenate the text_delta payloads an SSE event list carries."""
text = ""
for line in events:
for raw in line.split("\n"):
raw = raw.strip()
if not raw.startswith("data: "):
continue
data = json.loads(raw[len("data: ") :])
delta = data.get("delta", {})
if delta.get("type") == "text_delta":
text += delta.get("text", "")
return text
def test_anthropic_emitter_closes_reasoning_only_think_block():
# A reasoning-only reply streams <think>X live then shrinks to bare X at EOF.
# This emitter diffs cumulative snapshots and drops the shrink, so without a
# closing pass the client text would end on an unclosed <think>. finish()
# must balance it.
emitter = AnthropicStreamEmitter()
events = emitter.start("msg_1", "m")
events += emitter.feed({"type": "content", "text": "<think>The capital"})
events += emitter.feed({"type": "content", "text": "<think>The capital of France is Paris."})
# The generator's final bare-text shrink (dropped by the cumulative diff).
events += emitter.feed({"type": "content", "text": "The capital of France is Paris."})
events += emitter.finish()
assert _emitter_client_text(events) == "<think>The capital of France is Paris.</think>"
def test_anthropic_emitter_does_not_double_close_balanced_think():
# A reasoning-then-answer reply already closes its own </think>; the balancer
# must not append a second one.
emitter = AnthropicStreamEmitter()
events = emitter.start("msg_1", "m")
events += emitter.feed({"type": "content", "text": "<think>Thinking."})
events += emitter.feed({"type": "content", "text": "<think>Thinking.</think>Answer."})
events += emitter.finish()
assert _emitter_client_text(events) == "<think>Thinking.</think>Answer."
def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch):
import routes.inference as inf_mod
@ -889,6 +932,24 @@ class TestAnthropicToolNonStreaming:
assert tool_blocks[0]["name"] == "render_html"
assert tool_blocks[0]["input"] == {"code": "<!doctype html><html></html>"}
def test_display_strip_gates_on_declared_tools(self):
# A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared
# ``foo`` markup is prose and survives, the declared web_search rehearsal strips.
def _run_gen():
yield {
"type": "content",
"text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.',
}
tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
response = asyncio.run(
_anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools)
)
body = json.loads(response.body)
text = "".join(b["text"] for b in body["content"] if b["type"] == "text")
assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose
assert "web_search[ARGS]" not in text # active name stripped from display
# =====================================================================
# Pass-through emitter tests (client-side tool execution path)

194
studio/backend/tests/test_chat_eos_template_refresh.py Normal file
View file

@ -0,0 +1,194 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Mapper models whose own tokenizer ships no chat_template have their turn-end
eos resolved at LOAD from an empty template (document eos only). The effective
template is installed later, at generate time, via get_chat_template, so the
turn-end-eos cache must be refreshed then; otherwise generate_stream runs past
the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes).
"""
import sys
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parent.parent
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
# These tests construct InferenceBackend, pulling the full stack. CI may lack
# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup
# (RuntimeError); skip at module level so collection is not aborted (exit 2).
try:
from core.inference import inference as inf_mod # noqa: E402
from core.inference.inference import InferenceBackend # noqa: E402
except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent
pytest.skip(
f"full inference backend unavailable ({type(exc).__name__}: {exc})",
allow_module_level = True,
)
_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
_GEMMA = "{% for m in messages %}<start_of_turn>{{m.role}}\n{{m.content}}<end_of_turn>{% endfor %}"
class _FakeTokenizer:
def __init__(
self,
eos_id,
chat_template = "",
token_ids = None,
):
self.eos_token_id = eos_id
self.chat_template = chat_template
self.pad_token_id = eos_id
self.unk_token_id = None
self._ids = dict(token_ids or {})
def convert_tokens_to_ids(self, tok):
return self._ids.get(tok)
def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch):
import utils.datasets as ds
backend = InferenceBackend.__new__(InferenceBackend)
backend.active_model_name = "unsloth/qwen2.5-0.5b"
# No chat_template at load, so the cache stored only the document eos, though
# <|im_end|> is atomic in the vocab (unused until the mapper installs a template).
bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645})
model_info = {
"tokenizer": bare_tok,
"is_vision": False,
"chat_turn_end_eos_ids": [151643],
}
backend.models = {backend.active_model_name: model_info}
# The mapper installs a ChatML template (turns end with <|im_end|>) at generate time.
templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645})
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok)
monkeypatch.setattr(
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False
)
# Stub the tail so the generator runs through the refresh without a real model.
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
monkeypatch.setattr(
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
)
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
# After the template is applied the cache must include the ChatML turn-end id.
assert model_info["chat_turn_end_eos_ids"] == [151643, 151645]
def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch):
# Regression: get_chat_template can return a remapped tokenizer (Gemma: <end_of_turn>
# folded onto the eos id) while generate_stream re-reads the original. Resolving on
# the swap yields a narrower set, so the refresh must UNION, never overwrite.
import utils.datasets as ds
backend = InferenceBackend.__new__(InferenceBackend)
backend.active_model_name = "unsloth/gemma-2b-it"
# Original tokenizer (used by generate_stream): <end_of_turn>=107 distinct from
# eos=1, so the load-time cache resolved to [1, 107].
orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 107})
model_info = {
"tokenizer": orig_tok,
"is_vision": False,
"chat_turn_end_eos_ids": [1, 107],
}
backend.models = {backend.active_model_name: model_info}
# Destructively-swapped tokenizer: <end_of_turn> now maps onto eos id 1, so
# resolving on it yields only [1] (drops 107).
swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 1})
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok)
monkeypatch.setattr(
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False
)
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
monkeypatch.setattr(
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
)
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
# The load-time <end_of_turn>=107 must survive: overwriting with the swapped
# [1] would regress and loop past the turn.
assert model_info["chat_turn_end_eos_ids"] == [1, 107]
def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch):
# Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but
# get_chat_template folds it onto the doc-eos id. generate_stream uses the original,
# so read marker strings from the mapped template but ids from the original.
import utils.datasets as ds
backend = InferenceBackend.__new__(InferenceBackend)
backend.active_model_name = "01-ai/yi-6b"
# Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7.
orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
model_info = {
"tokenizer": orig_tok,
"is_vision": False,
"chat_turn_end_eos_ids": [2],
}
backend.models = {backend.active_model_name: model_info}
# Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2.
remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok)
monkeypatch.setattr(
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False
)
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
monkeypatch.setattr(
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
)
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
# The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2.
assert model_info["chat_turn_end_eos_ids"] == [2, 7]
class _FakeProcessor:
"""A ProcessorMixin-like container: carries the chat_template itself and
wraps the real text tokenizer as ``.tokenizer`` (the vision layout)."""
def __init__(self, chat_template, tokenizer):
self.chat_template = chat_template
self.tokenizer = tokenizer
def test_resolve_chat_eos_reads_vision_processor_template():
# Vision model: the chat_template lives on the processor while the inner tokenizer
# ships none. _resolve_chat_eos must read the marker from the processor but resolve
# its id on the inner tokenizer, and repair generation_config.
from types import SimpleNamespace
inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"<end_of_turn>": 107})
processor = _FakeProcessor(_GEMMA, inner_tok)
model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1))
backend = InferenceBackend.__new__(InferenceBackend)
backend.active_model_name = "unsloth/gemma-3-4b-it"
model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True}
backend.models = {backend.active_model_name: model_info}
backend._resolve_chat_eos(backend.active_model_name)
assert model_info["chat_turn_end_eos_ids"] == [1, 107]
# generation_config repaired so the vision .generate() path stops at the turn.
assert model.generation_config.eos_token_id == [1, 107]

11
studio/backend/tests/test_chat_history_routes.py
View file

@ -91,6 +91,17 @@ def test_chat_settings_payload_accepts_fast_mode_presets():
assert dumped["customPresets"][0]["params"]["fastMode"] is True
def test_chat_settings_payload_accepts_nudge_tool_calls():
# extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the
# frontend's persisted nudgeToolCalls needs a payload field (like
# autoHealToolCalls).
payload = chat_history.ChatSettingsPayload.model_validate(
{"autoHealToolCalls": True, "nudgeToolCalls": False}
)
dumped = payload.model_dump(exclude_unset = True)
assert dumped == {"autoHealToolCalls": True, "nudgeToolCalls": False}
def test_chat_inference_settings_covers_frontend_persisted_fields():
# Drift guard: every InferenceParams field the UI persists (all but
# checkpoint) must exist on ChatInferenceSettings, else extra="forbid"

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

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

157
studio/backend/tests/test_chat_template_tool_arguments.py Normal file
View file

@ -0,0 +1,157 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""apply_chat_template_for_generation must coerce assistant tool_call arguments
from the OpenAI JSON-string form to a dict before rendering. Strict tool
templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and
raise "Can only get item pairs from a mapping." on the string form when a prior
tool call is re-rendered on the next turn (MLX + transformers paths).
"""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parent.parent
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from core.inference.chat_template_helpers import ( # noqa: E402
_normalize_tool_call_arguments,
apply_chat_template_for_generation,
)
def _conv(arguments):
return [
{"role": "user", "content": "weather?"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"type": "function",
"id": "c1",
"function": {"name": "web_search", "arguments": arguments},
}
],
},
{"role": "tool", "name": "web_search", "content": "21C sunny"},
]
class _StrictTemplateTokenizer:
"""Mimics a strict Qwen tool template: rejects string tool_call arguments."""
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kw,
):
for msg in messages:
for call in msg.get("tool_calls", []) or []:
args = call.get("function", {}).get("arguments")
if isinstance(args, str):
raise TypeError("Can only get item pairs from a mapping.")
return "RENDERED"
def test_string_arguments_are_parsed_to_dict():
out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}'))
args = out[1]["tool_calls"][0]["function"]["arguments"]
assert args == {"query": "sweden"}
def test_dict_arguments_untouched_and_no_copy():
conv = _conv({"query": "sweden"})
assert _normalize_tool_call_arguments(conv) is conv
def test_non_json_string_left_as_is():
out = _normalize_tool_call_arguments(_conv("not json"))
assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json"
def test_render_succeeds_on_strict_template_with_string_arguments():
# Regression: strict template + string args used to raise.
result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}'))
assert result == "RENDERED"
class _RecordingTokenizer:
"""Lenient template: renders whatever arguments it is given (string or dict)."""
def __init__(self):
self.seen_arguments = None
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kw,
):
for msg in messages:
for call in msg.get("tool_calls", []) or []:
self.seen_arguments = call.get("function", {}).get("arguments")
return "RENDERED"
def test_lenient_template_receives_original_string_untouched():
# Lenient template must see the exact original string, not a coerced dict.
tok = _RecordingTokenizer()
apply_chat_template_for_generation(tok, _conv('{"query": "x"}'))
assert tok.seen_arguments == '{"query": "x"}'
def test_messages_without_tool_calls_pass_through_unchanged():
conv = [{"role": "user", "content": "hi"}]
assert _normalize_tool_call_arguments(conv) is conv
class _RaiseExceptionTemplateTokenizer:
"""Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via
``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError."""
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kw,
):
for msg in messages:
for call in msg.get("tool_calls", []) or []:
args = call.get("function", {}).get("arguments")
if isinstance(args, str):
raise ValueError(
"chat_template: tool_calls[].function.arguments must be a "
"JSON object (mapping), not a string."
)
return "RENDERED"
def test_render_succeeds_on_raise_exception_template_with_string_arguments():
# Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce.
result = apply_chat_template_for_generation(
_RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}')
)
assert result == "RENDERED"
def test_unrelated_template_error_still_propagates_with_dict_args():
# Failure unrelated to string args (dict args, nothing to coerce) must propagate.
class _AlwaysRaises:
def apply_chat_template(self, messages, **kw):
raise ValueError("template is broken")
with pytest.raises(ValueError, match = "broken"):
apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"}))

150
studio/backend/tests/test_chat_turn_end_eos.py Normal file
View file

@ -0,0 +1,150 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and
repair generation_config so a chat model whose eos is a bare document terminator
(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn
boundary instead of running past it and looping. Dependency-light: imported here
without the full inference stack.
"""
from __future__ import annotations
import sys
from pathlib import Path
_BACKEND = Path(__file__).resolve().parent.parent
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from core.inference.chat_eos import ( # noqa: E402
chat_eos_repair,
resolve_chat_turn_end_eos_ids,
resolve_chat_turn_end_eos_ids_using,
)
class _FakeTokenizer:
def __init__(
self,
eos_id,
chat_template = "",
token_ids = None,
unk_token_id = None,
):
self.eos_token_id = eos_id
self.chat_template = chat_template
self.unk_token_id = unk_token_id
self._ids = dict(token_ids or {})
def convert_tokens_to_ids(self, tok):
return self._ids.get(tok, self.unk_token_id)
# ---- resolve_chat_turn_end_eos_ids ---------------------------------------
_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
def test_qwen35_adds_im_end_from_template():
# eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046).
tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046})
assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046]
def test_marker_in_vocab_but_not_in_template_is_ignored():
# Base/coder model: <|im_end|> is in the vocab but the template does not use
# it, so it must not become a stop token.
tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046})
assert resolve_chat_turn_end_eos_ids(tok) == [248044]
def test_harmony_template_is_left_untouched():
# gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end.
harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"
tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007})
assert resolve_chat_turn_end_eos_ids(tok) == [200002]
def test_llama3_eot_id_from_template():
tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009})
assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009]
def test_gemma4_turn_marker_from_template():
# Gemma-4 ends turns with <turn|> while keeping a document eos, so <turn|> must
# be added as a stop token.
tok = _FakeTokenizer(
1, chat_template = "...<start_of_turn>...<turn|>...", token_ids = {"<turn|>": 106}
)
assert resolve_chat_turn_end_eos_ids(tok) == [1, 106]
def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer():
# map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id,
# but the original keeps it atomic. Reading marker STRINGS from the template but
# IDS on the original recovers the real turn-end id (7), not the doc-eos id (2).
template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7]
# Same tokenizer for both reproduces the plain resolve (load-time behaviour).
assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2]
def test_list_eos_preserved():
tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
assert resolve_chat_turn_end_eos_ids(tok) == [1, 2]
def test_missing_marker_maps_to_unk_and_is_skipped():
tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0)
assert resolve_chat_turn_end_eos_ids(tok) == [7]
def test_starling_barred_end_of_turn_from_template():
# OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from
# Gemma's <end_of_turn>). eos synced to </s>=2, turn marker at 32000.
starling = "GPT4 Correct Assistant: hi<|end_of_turn|>"
tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000})
assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000]
def test_dict_chat_template_scans_all_variants():
# Hermes-3 style: chat_template is a {name: template} dict. Detection must scan
# every variant, not bail because the container is not a plain str.
tmpl = {"default": "{{ messages }}", "tool_use": _CHATML}
tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
def test_list_of_dicts_chat_template_scans_all_variants():
# tokenizer_config.json stores multi-templates as a list of {name, template}.
tmpl = [{"name": "default", "template": _CHATML}]
tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
def test_dict_harmony_template_left_untouched():
# A multi-variant container whose variant is harmony must still be left alone.
tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"}
tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007})
assert resolve_chat_turn_end_eos_ids(tok) == [200002]
# ---- chat_eos_repair ------------------------------------------------------
def test_repair_adds_missing_turn_end():
assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046]
def test_repair_from_missing_generation_config_eos():
assert chat_eos_repair(None, [248046]) == [248046]
def test_repair_noop_when_already_covered():
assert chat_eos_repair([248046, 248044], [248046]) is None
def test_repair_noop_when_no_turn_end_ids():
assert chat_eos_repair(248044, []) is None

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

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

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

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

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

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

312
studio/backend/tests/test_gemma_tool_parse_edge_cases.py
View file

@ -1,15 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Edge cases in Gemma-native tool-call parsing.
Covers two failure modes:
1. A bare (unquoted) string argument that contains a comma, e.g.
``location:New York, NY`` -- the comma must not be treated as the next
key boundary, or the whole call is dropped.
2. A tool-call marker that appears INSIDE another call's argument string is
data, not a real call, so it must not be promoted to a second tool call.
"""
"""Gemma-native tool-call parsing edge cases: commas inside bare string values,
and markers inside another call's argument data staying data."""
from __future__ import annotations
@ -21,7 +14,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference.tool_call_parser import parse_tool_calls_from_text
from core.inference.tool_call_parser import (
_gemma_parse_value,
parse_tool_calls_from_text,
)
from core.tool_healing import strip_tool_call_markup
def _args(call: dict) -> dict:
@ -40,14 +37,22 @@ def test_bare_string_argument_with_comma_is_kept():
def test_normal_multi_key_arguments_still_split():
calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}<tool_call|>')
assert len(calls) == 1, calls
# Numbers stay numeric, bare strings get quoted, an explicit quoted comma
# stays inside its value.
assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"}
def test_empty_bare_value_becomes_empty_string_not_dropped():
# An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call).
calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}<tool_call|>")
assert len(calls) == 1, calls
assert _args(calls[0]) == {"query": "", "unit": "celsius"}
only = parse_tool_calls_from_text("<|tool_call>call:get{q:}<tool_call|>")
assert len(only) == 1, only
assert _args(only[0]) == {"q": ""}
def test_bare_value_with_timestamps_after_comma_is_kept():
# A comma followed by digits-then-colon (a timestamp/ratio) is value text,
# not a new key, so the whole query must be preserved as one argument.
# A comma before digits-then-colon (timestamp/ratio) is value text, not a key.
calls = parse_tool_calls_from_text(
"<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}<tool_call|>"
)
@ -55,9 +60,16 @@ def test_bare_value_with_timestamps_after_comma_is_kept():
assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"}
def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept():
# The wrapper-less Gemma form (no <|tool_call> markers) goes through the
# _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE.
calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}")
assert len(calls) == 1, calls
assert calls[0]["function"]["name"] == "web_search"
assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"}
def test_marker_inside_json_argument_is_not_a_second_call():
# A python call whose `code` argument contains a Gemma marker string. The
# marker is data and must not execute as a second `terminal` call.
content = (
'<tool_call>{"name":"python","arguments":{"code":'
'"x = 1 # <|tool_call>call:terminal{command:ls}<tool_call|>"}}</tool_call>'
@ -75,8 +87,6 @@ def test_two_separate_gemma_calls_both_parse():
def test_mixed_format_calls_preserve_document_order():
# A Gemma-native call precedes a JSON-format call in the text; tools execute
# in returned order, so `create` must come before `read`.
content = (
"<|tool_call>call:create{path:a}<tool_call|> then "
'<tool_call>{"name":"read","arguments":{"path":"a"}}</tool_call>'
@ -86,8 +96,6 @@ def test_mixed_format_calls_preserve_document_order():
def test_json_marker_inside_gemma_argument_is_not_a_second_call():
# The reverse of the JSON-outer case: a JSON-style marker inside a Gemma
# call's quoted argument is code text, not a second `terminal` call.
content = (
'<|tool_call>call:python{code:<|"|>'
'print(<tool_call>{"name":"terminal","arguments":{"command":"ls"}}</tool_call>)'
@ -98,18 +106,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call():
def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call():
# An UNQUOTED Gemma value containing a literal marker: the outer object fails
# to normalize (the inner braces/marker break the JSON), but the inner marker
# is nested in the outer candidate span, so it must not be promoted to a
# standalone `terminal` call. The safe outcome is no executed tool call.
# An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer
# candidate span, so it must not be promoted to a standalone `terminal` call (no tool call).
content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}<tool_call|>}<tool_call|>"
calls = parse_tool_calls_from_text(content)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_bare_string_array_argument_is_quoted():
# Gemma may emit an array of bare strings without per-element quotes; they
# must be quoted so the call is not dropped.
calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}<tool_call|>")
assert len(calls) == 1, calls
assert _args(calls[0]) == {"labels": ["bug", "ui"]}
@ -123,8 +127,6 @@ def test_array_keeps_numbers_and_quoted_elements():
def test_array_of_objects_is_normalised():
# Arrays of objects are a common tool-schema shape; their (unquoted) keys and
# bare values must be normalised too, not left verbatim, or the call drops.
calls = parse_tool_calls_from_text(
"<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}<tool_call|>"
)
@ -138,9 +140,6 @@ def test_nested_array_elements_are_normalised():
def test_gemma_marker_inside_xml_parameter_is_not_a_second_call():
# An XML-style <function=...> call whose <parameter=code> value contains a
# Gemma marker: the marker is the parameter's data, not a separate terminal
# call, so only the python call must be returned.
content = (
"<tool_call><function=python><parameter=code>"
"x = 1 # <|tool_call>call:terminal{command:ls}<tool_call|>"
@ -159,3 +158,254 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call():
)
calls = parse_tool_calls_from_text(content)
assert [c["function"]["name"] for c in calls] == ["python"], calls
def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call():
# A literal <think> inside a completed call's arguments is argument data; both calls must parse.
text = '[TOOL_CALLS]a{"x":"literal <think> marker"} b[ARGS]{"y":2}'
calls = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in calls] == ["a", "b"], calls
def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal():
# A genuine reasoning block still hides its rehearsal while a real call after it parses.
text = '<think>web_search[ARGS]{"q":"draft"}</think>real[ARGS]{"q":"go"}'
calls = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in calls] == ["real"], calls
def test_wrapperless_nested_object_argument_is_parsed():
# skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare.
calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}")
assert len(calls) == 1
assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3}
def test_wrapperless_array_argument_is_parsed():
calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}")
assert len(calls) == 1
assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2}
def test_wrapperless_deeply_nested_object_and_array_are_preserved():
# The single-pass parser must keep multi-level nesting (objects inside
# objects, arrays inside arrays) intact, not flatten or drop it.
calls = parse_tool_calls_from_text(
"call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}"
)
assert len(calls) == 1
assert _args(calls[0]) == {
"loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}},
"tags": ["a", "b", ["c", "d"]],
"n": 3,
}
def test_gemma_parse_array_advances_on_stray_brace():
# Regression: a stray '}' / ']' / ',' where an array element is expected must
# not stall _gemma_parse_value at the same index (it looped forever before).
from core.inference.tool_call_parser import _gemma_parse_array
items, end, closed = _gemma_parse_array("[a,}]", 0)
assert end == 5 and closed is True # consumed through the closing ']'
assert items[0] == "a"
def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping():
# Parse keeps the quoted close marker as data; strip removes the whole span.
text = '<|tool_call>call:python{code:<|"|>print("<tool_call|>")<|"|>}<tool_call|>'
calls = parse_tool_calls_from_text(text)
assert len(calls) == 1, calls
assert _args(calls[0]) == {"code": 'print("<tool_call|>")'}
assert strip_tool_call_markup("before " + text + " after") == "before after"
assert strip_tool_call_markup("before " + text + " after", final = True) == "before after"
def test_nested_xml_in_malformed_gemma_call_does_not_execute():
# The failed Gemma candidate's span still covers its nested <function=>.
text = (
"<|tool_call>call:outer{code:<function=terminal><parameter=command>id"
"</parameter></function></tool_call>, broken:{x}}<tool_call|>"
)
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_unbalanced_gemma_call_with_xml_does_not_execute():
# Unclosed braces cover to EOF, so the trailing <function=> is excluded.
text = (
"<|tool_call>call:outer{code:<function=terminal>"
"<parameter=command>id</parameter></function>"
)
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_standalone_function_xml_still_parses():
text = "<function=terminal><parameter=command>id</parameter></function>"
calls = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in calls] == ["terminal"], calls
def test_xml_between_braces_and_close_marker_does_not_execute():
# Coverage runs to the close marker, so <function=> in the gap is data.
text = (
"<|tool_call>call:outer{broken:{x}}<function=terminal>"
"<parameter=command>id</parameter></function><tool_call|>"
)
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_balanced_inner_call_inside_unclosed_outer_does_not_execute():
text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}<tool_call|>"
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_strip_preserves_text_after_malformed_gemma_close():
# Junk before the close is a malformed span: strip through it, keep the tail.
text = "pre <|tool_call>call:t{a:1} note <tool_call|> post"
assert strip_tool_call_markup(text) == "pre post"
assert strip_tool_call_markup(text, final = True) == "pre post"
def test_malformed_closed_gemma_span_is_stripped():
assert (
strip_tool_call_markup('before <|tool_call>{"name":"x"}<tool_call|> after')
== "before after"
)
def test_valid_call_after_missing_close_is_recovered():
# A close-less call covers only its braces, so the later call is recovered.
text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}<tool_call|>"
names_inc = [
c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True)
]
assert "b" in names_inc, names_inc
names_strict = [
c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False)
]
assert names_strict == ["b"], names_strict
def test_strip_non_final_keeps_incomplete_gemma_block():
text = "before <|tool_call>call:t{"
assert strip_tool_call_markup(text) == text
assert strip_tool_call_markup(text, final = True) == "before"
def test_json_call_between_gemma_braces_and_close_does_not_execute():
# A JSON call between the outer's braces and its close is covered data.
text = (
"<|tool_call>call:outer{broken:{x}}"
'<tool_call>{"name":"terminal","arguments":{"command":"id"}}</tool_call>'
"<tool_call|>"
)
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_gemma_call_between_gemma_braces_and_close_does_not_execute():
# Same escape with a Gemma-native inner marker.
text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}<tool_call|><tool_call|>"
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert "terminal" not in [c["function"]["name"] for c in calls], calls
def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener():
# The to-EOF Gemma sweep must not eat visible text after </function>.
text = (
'before <function=python><parameter=code>print("<|tool_call>")</parameter></function> after'
)
assert strip_tool_call_markup(text, final = True) == "before after"
assert strip_tool_call_markup(text) == "before after"
def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener():
# A call-form Gemma opener quoted in a closed block must not truncate it.
xml = "<function=python><parameter=code><|tool_call>call:t{</parameter></function>"
json_block = (
'<tool_call>{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}</tool_call>'
)
for block in (xml, json_block):
text = "before " + block + " after"
assert strip_tool_call_markup(text, final = True) == "before after", block
assert strip_tool_call_markup(text) == "before after", block
def test_function_sibling_after_close_less_gemma_marker_is_recovered():
# The close-less marker covers only its braces; the XML sibling is recovered.
text = (
"<|tool_call>call:bad{broken:{x}} "
"<function=terminal><parameter=command>id</parameter></function>"
)
for allow_incomplete in (True, False):
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
assert [c["function"]["name"] for c in calls] == ["terminal"], calls
def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered():
# A close token quoted in the later call must not extend the earlier
# close-less marker's coverage over that call.
gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|></tool_call><|"|>}<tool_call|>'
names = [
c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False)
]
assert names == ["b"], names
json_text = (
'<tool_call>{"name":"a","arguments":{}} '
'<tool_call>{"name":"b","arguments":{"x":"</tool_call>"}}</tool_call>'
)
names_j = [
c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False)
]
assert "b" in names_j, names_j
def test_gemma_parse_value_always_advances_on_stray_delimiter():
# A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the
# index by at least one, or a caller looping on it spins forever at 100% CPU (DoS).
for delim in (",", "}", "]"):
text = delim + "rest"
value, nxt, _explicit = _gemma_parse_value(text, 0)
assert nxt > 0, (delim, value, nxt)
def test_malformed_gemma_array_does_not_hang():
# ``[},]`` puts a stray ``}`` at the primitive position inside a list body.
# On the buggy parser this hangs the server; guard with a wall-clock timeout
# so the regression fails loudly instead of blocking CI forever.
import threading
result: dict = {}
def _run():
result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}<tool_call|>")
t = threading.Thread(target = _run, daemon = True)
t.start()
t.join(timeout = 10.0)
assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input"
def test_malformed_gemma_mapping_value_does_not_hang():
# A stray ``}`` where a mapping value is expected must also terminate.
import threading
result: dict = {}
def _run():
result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}<tool_call|>")
t = threading.Thread(target = _run, daemon = True)
t.start()
t.join(timeout = 10.0)
assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input"

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

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

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

File diff suppressed because it is too large Load diff

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

@ -587,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names():
import re as _re
from pathlib import Path
from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text()
m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL)
assert m, "could not extract _TOOL_XML_RE"
ns: dict = {"_re": _re}
ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC}
exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns)
rx = ns["_TOOL_XML_RE"]
stripped = rx.sub(

51
studio/backend/tests/test_mlx_inference_backend.py
View file

@ -100,6 +100,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):
]
assert backend._is_vlm is False
assert isinstance(backend._tokenizer, _DummyTokenizer)
# Non-LoRA text model: no base_model on the record.
assert backend.models["fake/text"]["base_model"] is None
def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch):
# A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template
# lives on the base model.
_install_fake_mlx(monkeypatch)
calls = []
_install_fake_fast_mlx(monkeypatch, calls)
from core.inference.mlx_inference import MLXInferenceBackend
backend = MLXInferenceBackend()
config = SimpleNamespace(
identifier = "fake/text-adapter",
is_vision = False,
is_lora = True,
base_model = "fake/text-base",
)
assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token")
record = backend.models["fake/text-adapter"]
assert record["is_lora"] is True
assert record["base_model"] == "fake/text-base"
def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite(
@ -188,12 +214,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import MLXInferenceBackend
captured = {}
# The text path renders once with tools, then the native-template fallback makes a second no-
# tools probe call (tools=None) to detect whether the template dropped the schema.
captured_calls = []
def _fake_apply(tokenizer, messages, **kwargs):
captured["tokenizer"] = tokenizer
captured["messages"] = messages
captured["kwargs"] = kwargs
captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs})
return "<rendered prompt>"
monkeypatch.setattr(
@ -248,8 +274,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
)
)
assert out == ["hi"]
# The toggled kwargs must reach the chat-template helper.
assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}]
assert captured["kwargs"]["enable_thinking"] is True
assert captured["kwargs"]["reasoning_effort"] == "medium"
assert captured["kwargs"]["preserve_thinking"] is True
# The toggled kwargs must reach the chat-template helper on the real render
# (one of the calls carries the tools; the fallback probe passes tools=None).
tool_renders = [
c
for c in captured_calls
if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}]
]
assert tool_renders, captured_calls
render = tool_renders[0]
assert render["kwargs"]["enable_thinking"] is True
assert render["kwargs"]["reasoning_effort"] == "medium"
assert render["kwargs"]["preserve_thinking"] is True

23
studio/backend/tests/test_mlx_repair.py
View file

@ -271,6 +271,29 @@ def test_stack_available_requires_runtime_imports_and_versions(monkeypatch):
assert imported == list(mr._MLX_RUNTIME_IMPORTS)
def test_mlx_packages_exclude_known_bad_mlx_lm():
# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5); the install spec
# must exclude it so the resolver picks 0.31.2 or >=0.31.4. See mlx-lm #1242.
(mlx_lm_spec,) = [p for p in mr.MLX_PACKAGES if p.startswith("mlx-lm")]
assert mlx_lm_spec == "mlx-lm>=0.22.0,!=0.31.3"
@pytest.mark.parametrize("bad_form", ["0.31.3", "0.31.3.0"])
def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form):
# An installed 0.31.3 counts as unsatisfied so the self-heal replaces it;
# parsed-Version compare also catches the trailing-zero form 0.31.3.0.
import importlib.metadata as metadata
def _version(name):
return bad_form if name == "mlx-lm" else mr._MLX_MIN_VERSIONS[name]
monkeypatch.setattr(metadata, "version", _version)
monkeypatch.setattr(
mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports")
)
assert mr.mlx_stack_available() is False
def test_no_op_off_apple_silicon(monkeypatch):
monkeypatch.setattr(mr, "is_apple_silicon", lambda: False)
called = {"n": 0}

176
studio/backend/tests/test_native_template_trust_remote_code.py Normal file
View file

@ -0,0 +1,176 @@
# 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 trust_remote_code in the native-template fallback.
``render_native_template`` re-fetches a model's native chat template from its
repo when an Unsloth override template (mistral, gemma-4) dropped the tools
schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo
carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use
that same consent or transformers raises (it requires ``trust_remote_code`` to
instantiate a custom tokenizer class), the ``except`` swallows it, and the
request silently keeps the tool-dropping prompt even though the user already
consented to remote code for the model load.
These tests pin that the stored ``trust_remote_code`` is threaded to the reload,
that the reload is skipped (returns ``None`` without executing code) when no
consent is stored, and that both backend ``model_info`` dicts persist the flag at
load time so the read lands on a value ``load_model`` actually set.
"""
from __future__ import annotations
import importlib.util
import sys
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)
# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the
# transformers import deferred inside the function). Load it directly so the test
# runs without importing the heavy ``core.inference`` package (unsloth / torch).
_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py"
_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH)
chat_template_helpers = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(chat_template_helpers)
render_native_template = chat_template_helpers.render_native_template
# A native template that emits a tools section only when tools are provided, so the
# with-tools vs no-tools render differs and ``render_native_template`` accepts it.
_NATIVE_TEMPLATE = (
"{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}"
"{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}"
"{% if add_generation_prompt %}assistant:{% endif %}"
)
_MESSAGES = [{"role": "user", "content": "what is the weather"}]
_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}]
class _JinjaTokenizer:
"""Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``.
Stands in for the live model tokenizer that ``render_native_template`` shallow-
copies and re-points at the native template before rendering.
"""
def __init__(self, chat_template):
self.chat_template = chat_template
def apply_chat_template(
self,
messages,
tokenize = False,
add_generation_prompt = True,
tools = None,
**kwargs,
):
from jinja2 import BaseLoader, Environment
env = Environment(loader = BaseLoader())
return env.from_string(self.chat_template).render(
messages = messages,
tools = tools,
add_generation_prompt = add_generation_prompt,
)
def _install_custom_code_tokenizer(monkeypatch):
"""Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise
unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the
native template. Records the ``trust_remote_code`` it was called with."""
pytest.importorskip("jinja2")
from transformers import AutoTokenizer
calls = {}
def fake_from_pretrained(
model_id,
*args,
trust_remote_code = False,
token = None,
**kwargs,
):
calls["trust_remote_code"] = trust_remote_code
calls["model_id"] = model_id
calls["token"] = token
if not trust_remote_code:
# Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code:
# has_remote_code and not has_local_code and not trust_remote_code -> ValueError.
raise ValueError(
f"The repository {model_id} contains custom code which must be executed "
"to correctly load the model. Please pass the argument "
"`trust_remote_code=True` to allow custom code to be run."
)
return _JinjaTokenizer(_NATIVE_TEMPLATE)
monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained))
return calls
def _model_info(trust_remote_code):
return {
"native_chat_template": None, # force the repo reload path
"base_model": None, # non-LoRA: template_source == active_model_name
"trust_remote_code": trust_remote_code,
# Live tokenizer that gets shallow-copied + re-pointed at the native template.
"tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"),
}
def test_native_reload_passes_stored_trust_remote_code(monkeypatch):
"""With ``trust_remote_code`` stored on ``model_info`` the custom-code reload
succeeds and the tools-advertising native prompt is returned. This FAILS before
the fix (reload omits the flag, raises, is swallowed, returns None)."""
calls = _install_custom_code_tokenizer(monkeypatch)
model_info = _model_info(trust_remote_code = True)
out = render_native_template(
model_info = model_info,
active_model_name = "acme/custom-tokenizer-model",
messages = _MESSAGES,
tools = _TOOLS,
)
assert out is not None, "native fallback should render the tools prompt with consent"
assert "[AVAILABLE_TOOLS]" in out
assert "get_weather" in out
assert calls["trust_remote_code"] is True # the stored consent was threaded through
# A successful fetch is cached so the next tool turn skips the reload.
assert model_info["native_chat_template"] == _NATIVE_TEMPLATE
def test_native_reload_without_consent_returns_none(monkeypatch):
"""Without stored consent the custom-code reload raises, is swallowed, and
``render_native_template`` returns None (no unconsented code execution). Proves
the stored flag -- not a hard-coded True -- drives the reload."""
calls = _install_custom_code_tokenizer(monkeypatch)
model_info = _model_info(trust_remote_code = False)
out = render_native_template(
model_info = model_info,
active_model_name = "acme/custom-tokenizer-model",
messages = _MESSAGES,
tools = _TOOLS,
)
assert out is None
assert calls["trust_remote_code"] is False
# A failed fetch must not be cached as "no template" (would pin the tool drop).
assert model_info["native_chat_template"] is None
def test_backend_model_info_persists_trust_remote_code():
"""Both backends must store ``trust_remote_code`` on their per-model info dict so
``render_native_template`` can source the consent value. Guards against the read
landing on a key ``load_model`` never sets (which would silently no-op the fix)."""
inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text()
mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text()
assert '"trust_remote_code": trust_remote_code,' in inf
assert '"trust_remote_code": trust_remote_code,' in mlx

99
studio/backend/tests/test_nudge_tool_calls_wiring.py Normal file
View file

@ -0,0 +1,99 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy.
Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths
(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat +
Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off).
Mechanism (verified here without loading a model):
* every backend tool-loop entry point accepts and forwards ``nudge_tool_calls``
(safetensors -> ``InferenceBackend``; MLX -> ``InferenceOrchestrator``; both
call the shared ``run_safetensors_tool_loop``; GGUF -> ``LlamaCppBackend``);
* the safetensors/MLX loop gates the retry on a truthy flag (new retry ->
opt-in), while the GGUF loop keeps its pre-existing default-on behaviour
(``None`` keeps nudging) so an omitted flag never disables GGUF;
* the API request models default the flag to ``None`` (opt-in / off);
* the Studio-facing routes forward the request's flag, and the Studio frontend
sends ``nudge_tool_calls: true`` -- exercised behaviourally in
``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``.
"""
import inspect
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.orchestrator import InferenceOrchestrator
from core.inference.safetensors_agentic import run_safetensors_tool_loop
try:
# core.inference.inference imports unsloth at module scope, which requires
# unsloth_zoo. The dependency-light backend CI matrix job does not install
# it, so the safetensors InferenceBackend is folded into the checks below
# only when the unsloth stack is importable (local runs / full CI); the
# other entry points are always checked.
from core.inference.inference import InferenceBackend
except ImportError:
InferenceBackend = None
def _params(fn):
return inspect.signature(fn).parameters
def test_shared_loop_accepts_nudge_flag():
assert "nudge_tool_calls" in _params(run_safetensors_tool_loop)
def test_backends_accept_the_flag():
methods = [
InferenceOrchestrator.generate_chat_completion_with_tools,
LlamaCppBackend.generate_chat_completion_with_tools,
]
if InferenceBackend is not None: # safetensors path; needs the unsloth stack
methods.append(InferenceBackend.generate_chat_completion_with_tools)
for method in methods:
assert "nudge_tool_calls" in _params(method), method.__qualname__
def test_delegating_backends_forward_the_flag_to_the_shared_loop():
# safetensors (in-process transformers) and MLX (parent-process orchestrator)
# both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop
# and consumes the flag directly (asserted separately by the gate test).
methods = [InferenceOrchestrator.generate_chat_completion_with_tools]
if InferenceBackend is not None: # safetensors path; needs the unsloth stack
methods.append(InferenceBackend.generate_chat_completion_with_tools)
for method in methods:
src = inspect.getsource(method)
assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__
def test_safetensors_loop_is_opt_in_while_gguf_stays_default_on():
# Safetensors/MLX: the retry is new here, so it requires a truthy flag.
sf_src = inspect.getsource(run_safetensors_tool_loop)
assert "and nudge_tool_calls" in sf_src
# GGUF: pre-existing nudge must not be accidentally disabled -- an omitted
# (None) flag keeps nudging; only an explicit False turns it off.
gguf_src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools)
assert "nudge_tool_calls is None or nudge_tool_calls" in gguf_src
def test_api_request_models_default_the_flag_off():
from models.inference import AnthropicMessagesRequest, ChatCompletionRequest
for model in (ChatCompletionRequest, AnthropicMessagesRequest):
field = model.model_fields["nudge_tool_calls"]
assert field.default is None, model.__name__
def test_studio_routes_forward_the_request_flag():
# The Studio chat frontend posts to /v1/chat/completions and /v1/messages
# with nudge_tool_calls=true; the route handlers forward the request value
# (external API clients that omit it fall back to the opt-in default).
from routes import inference as routes_inference
for handler in (
routes_inference.openai_chat_completions,
routes_inference.anthropic_messages,
):
src = inspect.getsource(handler)
assert "nudge_tool_calls = payload.nudge_tool_calls" in src, handler.__name__

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

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

1622
studio/backend/tests/test_orchestrator_unload_cancel.py Normal file
View file

File diff suppressed because it is too large Load diff

89
studio/backend/tests/test_passthrough_healing.py
View file

@ -280,6 +280,56 @@ class TestStreamHealer:
assert [c["id"] for c in calls] == ["call_0", "call_1"]
assert _events_text(events).strip() == "then"
def test_mistral_array_multiple_calls_all_promoted_in_stream(self):
# A canonical Mistral [TOOL_CALLS] array carries several calls under a
# SINGLE signal. Draining only the first call would leave the residue
# starting at ",{...}]" (no signal), so later calls in the same array
# must be promoted in the same pass, not flushed as raw text.
healer = StreamToolCallHealer({"get_weather", "get_time"})
array = (
'[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},'
'{"name":"get_time","arguments":{"tz":"UTC"}}]'
)
events = healer.feed(array) + healer.finalize()
calls = _events_calls(events)
assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"]
assert [c["id"] for c in calls] == ["call_0", "call_1"]
assert _events_text(events) == ""
def test_mistral_array_multiple_calls_promoted_char_by_char(self):
healer = StreamToolCallHealer({"get_weather", "get_time"})
array = (
'[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},'
'{"name":"get_time","arguments":{"tz":"UTC"}}]'
)
events = []
for ch in array:
events += healer.feed(ch)
events += healer.finalize()
calls = _events_calls(events)
assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"]
assert _events_text(events) == ""
def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self):
# A mid-array element for a tool that is not declared must survive as
# text while the declared neighbours on either side still promote in
# document order.
healer = StreamToolCallHealer({"a", "c"})
array = (
'[TOOL_CALLS][{"name":"a","arguments":{}},'
'{"name":"b","arguments":{}},{"name":"c","arguments":{}}]'
)
events = healer.feed(array) + healer.finalize()
assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"]
assert '"b"' in _events_text(events)
def test_mistral_array_then_trailing_prose(self):
healer = StreamToolCallHealer({"a", "b"})
array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]'
events = healer.feed(f"{array} all done") + healer.finalize()
assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"]
assert "all done" in _events_text(events)
def test_incomplete_call_healed_at_finalize(self):
healer = StreamToolCallHealer({"Bash"})
events = healer.feed('<tool_call>{"name":"Bash","arguments":{"cmd":"ls"}}')
@ -1356,3 +1406,42 @@ class TestOpenaiStreamingRoute:
assert chunks[0] == line + "\n\n" # byte-for-byte relay
asyncio.run(_run())
class TestHealerSignalAlignment:
"""The passthrough healer buffers only formats its parser can promote.
The loops' bare [ARGS] rehearsal signal is gated on active tool names
there; ungated in the healer it would stall legitimate prose until
finalization without ever producing a promotable call."""
def test_heal_signals_are_promotable_formats_only(self):
from core.inference.passthrough_healing import _HEAL_SIGNALS
assert set(_HEAL_SIGNALS) == {
"<tool_call>",
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
}
def test_prose_with_bare_args_marker_streams_through(self):
healer = StreamToolCallHealer({"Bash"})
chunks = [
"Use the pattern foo",
"[ARGS] in templates when calling tools, ",
"and remember to close it.",
]
streamed = ""
for chunk in chunks:
streamed += _events_text(healer.feed(chunk))
# Incremental relay: nothing withheld for finalize.
assert streamed == "".join(chunks)
final = healer.finalize()
assert not _events_calls(final)
assert not healer.healed
def test_bracket_tool_calls_still_promote_in_stream(self):
healer = StreamToolCallHealer({"web_search"})
events = healer.feed('[TOOL_CALLS]web_search{"query": "unsloth docs"}') + healer.finalize()
(call,) = _events_calls(events)
assert call["function"]["name"] == "web_search"
assert healer.healed

1011
studio/backend/tests/test_pr5624_regressions.py Normal file
View file

File diff suppressed because it is too large Load diff

252
studio/backend/tests/test_presence_penalty.py Normal file
View file

@ -0,0 +1,252 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""Presence-penalty parity between the GGUF path and the safetensors/MLX paths.
The safetensors path historically dropped ``presence_penalty``, so the SAME model
looked worse served as safetensors. These tests pin the processor semantics
(subtract once per distinct completion token, prompt excluded, presence not
frequency, zero a no-op, negatives raise) plus a param-propagation regression
over route -> orchestrator cmd -> worker gen_kwargs.
"""
import threading
import pytest
import torch
from core.inference.presence_penalty import (
apply_presence_penalty,
_make_presence_penalty_processor,
)
def test_seen_token_gets_exactly_minus_penalty_unseen_unchanged():
input_ids = torch.tensor([[0, 1, 3]]) # prompt [0, 1], completion [3]
scores = torch.zeros(1, 5)
out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 2)
assert out[0, 3].item() == pytest.approx(-1.5)
for tok in (0, 1, 2, 4):
assert out[0, tok].item() == pytest.approx(0.0)
def test_multiplicity_ignored_presence_not_frequency():
# Token 3 emitted three times -> still a single -penalty (presence, not freq).
input_ids = torch.tensor([[0, 3, 3, 3]])
scores = torch.zeros(1, 5)
out = apply_presence_penalty(input_ids, scores, penalty = 2.0, prompt_len = 1)
assert out[0, 3].item() == pytest.approx(-2.0)
def test_negative_penalty_raises_seen_logits():
input_ids = torch.tensor([[0, 2]])
scores = torch.zeros(1, 4)
out = apply_presence_penalty(input_ids, scores, penalty = -0.5, prompt_len = 1)
assert out[0, 2].item() == pytest.approx(0.5)
def test_prompt_tokens_excluded():
# Token 7 is prompt-only (untouched); token 4 in the completion is penalized.
input_ids = torch.tensor([[7, 4, 4]])
scores = torch.zeros(1, 8)
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
assert out[0, 7].item() == pytest.approx(0.0)
assert out[0, 4].item() == pytest.approx(-1.0)
def test_batch_rows_isolated():
input_ids = torch.tensor([[0, 1], [0, 2]]) # row completions [1] and [2]
scores = torch.zeros(2, 4)
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
assert out[0, 1].item() == pytest.approx(-1.0)
assert out[0, 2].item() == pytest.approx(0.0)
assert out[1, 2].item() == pytest.approx(-1.0)
assert out[1, 1].item() == pytest.approx(0.0)
def test_zero_penalty_is_noop():
input_ids = torch.tensor([[0, 1, 2]])
scores = torch.randn(1, 5)
original = scores.clone()
out = apply_presence_penalty(input_ids, scores, penalty = 0.0, prompt_len = 1)
assert torch.equal(out, original)
def test_empty_completion_is_noop():
# prompt_len covers the whole sequence -> nothing generated yet.
input_ids = torch.tensor([[0, 1, 2]])
scores = torch.randn(1, 5)
original = scores.clone()
out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 3)
assert torch.equal(out, original)
def test_out_of_vocab_id_ignored():
# A generated id >= vocab_size (defensive) must not index out of bounds.
input_ids = torch.tensor([[0, 9]])
scores = torch.zeros(1, 5) # vocab 5, token 9 is out of range
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
assert torch.equal(out, torch.zeros(1, 5))
def test_negative_generated_id_ignored():
# A negative generated id (defensive) must be dropped, not wrap to scores[-1].
input_ids = torch.tensor([[0, -1]])
scores = torch.zeros(1, 5)
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
# Nothing penalized; in particular the last row (the numpy/torch wrap target
# for id -1) is untouched.
assert torch.equal(out, torch.zeros(1, 5))
def test_mixed_oob_negative_and_valid_ids_only_in_range_penalized():
# Completion mixes a valid id (1), an out-of-vocab id (9 >= vocab 5) and a
# negative id (-1). Only the in-range distinct id is penalized; OOB/negative
# ids are ignored with no crash and no wrong-index wrap. This fails under the
# old ``seen[seen < vocab_size]`` filter (id -1 wraps to the last row) and
# passes only with the both-ends bound.
input_ids = torch.tensor([[0, 1, 9, -1, 1]]) # prompt [0], completion [1, 9, -1, 1]
scores = torch.zeros(1, 5)
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
expected = torch.zeros(1, 5)
expected[0, 1] = -1.0 # once per distinct in-range id (multiplicity ignored)
assert torch.equal(out, expected)
assert out[0, 4].item() == pytest.approx(0.0) # id -1 did not wrap to the last row
def test_dtype_and_device_preserved():
input_ids = torch.tensor([[0, 1]])
scores = torch.zeros(1, 4, dtype = torch.float16)
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
assert out.dtype == torch.float16
assert out.device == scores.device
def test_processor_none_when_zero():
assert _make_presence_penalty_processor(0.0, prompt_len = 0) is None
def test_processor_applies_penalty():
proc = _make_presence_penalty_processor(1.5, prompt_len = 2)
assert proc is not None
input_ids = torch.tensor([[0, 1, 3]])
scores = torch.zeros(1, 5)
out = proc(input_ids, scores)
assert out[0, 3].item() == pytest.approx(-1.5)
def test_processor_composes_with_other_processors():
# LogitsProcessorList must run our processor alongside a pre-existing one.
from transformers import LogitsProcessor, LogitsProcessorList
class _AddToTokenZero(LogitsProcessor):
def __call__(self, input_ids, scores):
scores[:, 0] = scores[:, 0] + 100.0
return scores
presence = _make_presence_penalty_processor(1.0, prompt_len = 1)
combined = LogitsProcessorList([_AddToTokenZero(), *presence])
input_ids = torch.tensor([[5, 2]]) # completion = [2]
scores = torch.zeros(1, 6)
out = combined(input_ids, scores)
assert out[0, 0].item() == pytest.approx(100.0) # other processor ran
assert out[0, 2].item() == pytest.approx(-1.0) # presence ran
def test_mlx_presence_penalty_callable():
mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS")
from core.inference.mlx_inference import _make_mlx_presence_penalty_processor
proc = _make_mlx_presence_penalty_processor(1.5)
# First call = prompt only (latches prompt_len, penalizes nothing).
prompt = mx.array([10, 11])
logits0 = mx.zeros((1, 20))
out0 = proc(prompt, logits0)
assert float(out0[0, 10]) == pytest.approx(0.0)
# Second call: one completion token (5) appended -> penalized once.
seq = mx.array([10, 11, 5])
logits1 = mx.zeros((1, 20))
out1 = proc(seq, logits1)
assert float(out1[0, 5]) == pytest.approx(-1.5)
assert float(out1[0, 10]) == pytest.approx(0.0) # prompt token untouched
def test_mlx_presence_penalty_bounds_out_of_range_ids():
# Documents (and, on Apple Silicon CI, enforces) the intended MLX bound:
# out-of-vocab and negative completion ids must be ignored. MLX does no
# bounds checking and OOB indexing is undefined behavior (crash / memory
# corruption), so the processor routes stray ids to a discarded scratch slot
# and penalizes only in-range distinct ids -- matching the torch filter
# seen[(seen >= 0) & (seen < vocab)]. Skips off arm64 macOS where MLX is absent.
mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS")
from core.inference.mlx_inference import _make_mlx_presence_penalty_processor
proc = _make_mlx_presence_penalty_processor(1.0)
proc(mx.array([10, 11]), mx.zeros((1, 8))) # first call latches prompt_len = 2
# Completion appends a valid id (3), an out-of-vocab id (99 >= vocab 8) and a
# negative id (-1); only the in-range id is penalized and nothing crashes.
seq = mx.array([10, 11, 3, 99, -1])
out = proc(seq, mx.zeros((1, 8)))
assert float(out[0, 3]) == pytest.approx(-1.0)
for tok in range(8):
if tok != 3:
assert float(out[0, tok]) == pytest.approx(0.0)
# Param propagation: route payload -> orchestrator cmd -> worker gen_kwargs
_SAMPLING = {
"temperature": 0.7,
"top_p": 0.8,
"top_k": 20,
"min_p": 0.05,
"repetition_penalty": 1.1,
"presence_penalty": 1.5,
}
def test_orchestrator_cmd_carries_all_sampling_params():
from core.inference.orchestrator import InferenceOrchestrator
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
cmd = o._build_generate_cmd(
"req1",
None,
messages = [{"role": "user", "content": "hi"}],
max_new_tokens = 128,
**_SAMPLING,
)
for key, val in _SAMPLING.items():
assert cmd[key] == val, f"{key} dropped/altered in orchestrator cmd"
def test_worker_forwards_all_sampling_params_to_backend():
from core.inference.worker import _handle_generate
class _RecordingBackend:
last_generation_stats = None
def __init__(self):
self.received = None
def generate_chat_response(self, **kwargs):
self.received = kwargs
return iter(()) # empty stream -> loop exits, gen_done is sent
class _FakeQueue:
def __init__(self):
self.items = []
def put(self, item):
self.items.append(item)
cmd = {
"type": "generate",
"request_id": "r",
"messages": [{"role": "user", "content": "hi"}],
"max_new_tokens": 128,
**_SAMPLING,
}
backend = _RecordingBackend()
_handle_generate(backend, cmd, _FakeQueue(), threading.Event())
assert backend.received is not None
for key, val in _SAMPLING.items():
assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs"

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

@ -59,6 +59,7 @@ from models.inference import (
ResponsesUsage,
)
from routes.inference import (
_ResponsesReasoningExtractor,
_SameTaskStreamingResponse,
_build_chat_request,
_chat_tool_calls_to_responses_output,
@ -795,6 +796,7 @@ class TestResponsesNonStreamingAdapter:
def test_monitor_records_translated_visible_text(self, monkeypatch):
import routes.inference as inf_mod
import routes.inference as inf_mod
async def fake_chat_completions(chat_req, request):
assert request.state.skip_api_monitor is True
@ -1988,6 +1990,123 @@ class TestTranslatedMessagesValidate:
ChatMessage(**m.model_dump(exclude_none = True))
# reasoning_prefilled: enable_thinking templates prefill an unclosed <think>, so
# generation begins inside the block; the extractor must start in reasoning.
class TestReasoningPrefilledExtractor:
def test_prefilled_single_feed_splits_lone_close(self):
# T1: reasoning...</think>answer with a prefilled (unseen) open tag.
reasoning, visible = _extract_responses_reasoning(
"plan</think>answer",
parse_think_markers = True,
reasoning_prefilled = True,
)
assert reasoning == "plan"
assert visible == "answer"
def test_prefilled_never_closed_is_all_reasoning(self):
# T2: truncated mid-thought (no </think>) -> all reasoning (GGUF parity).
reasoning, visible = _extract_responses_reasoning(
"still thinking with no close",
parse_think_markers = True,
reasoning_prefilled = True,
)
assert reasoning == "still thinking with no close"
assert visible == ""
def test_prefilled_close_split_across_feeds(self):
# T3: </think> straddles two feed() calls; holdback resolves it.
ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True)
r1, v1 = ex.feed("plan</th")
r2, v2 = ex.feed("ink>ans")
fr, fv = ex.finish()
assert (r1 + r2 + fr) == "plan"
assert (v1 + v2 + fv) == "ans"
def test_prefilled_close_split_one_char_per_feed(self):
# T4: every char in its own feed still splits correctly.
ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True)
reasoning, visible = "", ""
for ch in "plan</think>x":
r, v = ex.feed(ch)
reasoning += r
visible += v
fr, fv = ex.finish()
assert (reasoning + fr) == "plan"
assert (visible + fv) == "x"
def test_prefilled_empty_generation(self):
# T5: nothing generated.
reasoning, visible = _extract_responses_reasoning(
"",
parse_think_markers = True,
reasoning_prefilled = True,
)
assert reasoning == ""
assert visible == ""
def test_prefilled_whitespace_after_close_is_visible(self):
# T6: Qwen commonly emits </think>\n\n before the answer.
reasoning, visible = _extract_responses_reasoning(
"plan</think>\n\nanswer",
parse_think_markers = True,
reasoning_prefilled = True,
)
assert reasoning == "plan"
assert visible == "\n\nanswer"
def test_prefilled_stray_open_tag_is_suppressed(self):
# T7: a re-emitted literal <think> inside prefilled reasoning is dropped,
# not leaked into the drawer (covers enable_thinking_effort full-tag output).
reasoning, visible = _extract_responses_reasoning(
"a<think>b</think>c",
parse_think_markers = True,
reasoning_prefilled = True,
)
assert reasoning == "ab"
assert visible == "c"
assert "<think>" not in reasoning
def test_prefilled_close_at_start_empty_reasoning(self):
# T8: model closed immediately (empty reasoning) then answered.
reasoning, visible = _extract_responses_reasoning(
"</think>hi",
parse_think_markers = True,
reasoning_prefilled = True,
)
assert reasoning == ""
assert visible == "hi"
def test_not_prefilled_lone_close_preserves_current_behavior(self):
# T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard).
reasoning, visible = _extract_responses_reasoning(
"reasoning</think>ans",
parse_think_markers = True,
reasoning_prefilled = False,
)
assert reasoning == ""
assert visible == "reasoningans"
def test_not_prefilled_full_pair_still_splits(self):
# T10: normal explicit <think>..</think> (GGUF / Harmony) unchanged.
reasoning, visible = _extract_responses_reasoning(
"<think>r</think>v",
parse_think_markers = True,
reasoning_prefilled = False,
)
assert reasoning == "r"
assert visible == "v"
def test_prefilled_ignored_when_markers_not_parsed(self):
# T11: a non-reasoning model passes text through even with reasoning_prefilled False.
reasoning, visible = _extract_responses_reasoning(
"just an answer",
parse_think_markers = False,
reasoning_prefilled = False,
)
assert reasoning == ""
assert visible == "just an answer"
# =====================================================================
# Streaming passthrough healing — text-form calls promoted in order
# =====================================================================

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

@ -11,6 +11,8 @@ from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
_backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
@ -46,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }}
"""
# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort
# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders
# identically to thinking-on-without-the-preamble), so the literal scan alone
# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to
# expose the encoder's full none/high/max ladder.
DEEPSEEK_V4_TEMPLATE = (
"{%- if not thinking is defined %}"
"{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}"
"{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n"
"{%- if thinking and reasoning_effort == 'max' %}"
"{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n"
"{%- for message in messages %}{{- message.content }}{%- endfor %}"
)
PLAIN_TEMPLATE = """
{%- for message in messages %}
{{- message.role + ': ' + message.content + '\\n' }}
@ -88,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false():
assert flags["reasoning_style"] == "enable_thinking"
def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max():
"""DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble.
Classified as the hybrid style with the full none/high/max ladder even
though the template only branches on 'max'."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF")
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
assert flags["reasoning_always_on"] is False
def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected():
"""The 'high' injection is scoped to deepseek-v4: a different model whose
template only branches on 'max' keeps ['max'] (no phantom 'high')."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF")
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["max"]
def test_detect_safetensors_features_passes_template_through_to_classifier():
"""Route wrapper forwards a real template to the inner classifier."""
from routes.inference import _detect_safetensors_features
@ -127,9 +167,8 @@ def test_detect_safetensors_features_gptoss_disables_tools():
assert flags["supports_tools"] is False
# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS],
# which our parser can't read. The route helper must not flip supports_tools=True
# for them, else the UI enables a pill the agentic loop can't honour.
# Llama-3 / Mistral / Gemma 4 tool-call formats are now parser-supported, so supports_tools=True
# must hold for all of them; only templates matching none of the five known markers are suppressed.
LLAMA3_TEMPLATE = """
{%- if tools %}
@ -161,27 +200,188 @@ MISTRAL_TEMPLATE = """
{%- endfor %}
"""
GEMMA4_TEMPLATE = """
{%- if tools %}
{{- 'Tools available. Emit calls as ' }}
{{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}<tool_call|>' }}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
"""
def test_detect_safetensors_features_llama3_template_suppresses_tools():
"""Llama-3 emits <|python_tag|>; safetensors loop cannot parse it."""
def test_detect_safetensors_features_llama3_template_keeps_tools_on():
"""Llama-3 emits <|python_tag|>; parser now supports it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE)
assert flags["supports_tools"] is False
assert flags["supports_tools"] is True
def test_detect_safetensors_features_mistral_template_suppresses_tools():
"""Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it."""
def test_detect_safetensors_features_mistral_template_keeps_tools_on():
"""Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses
(the shared bracket-tag parser). The gate must no longer suppress it, or the
PR's Mistral tool support is unreachable through normal capability detection."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3")
flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE)
assert flags["supports_tools"] is True
def test_detect_safetensors_features_gemma4_template_keeps_tools_on():
"""Gemma 4 emits <|tool_call>; parser now supports it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit")
flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE)
assert flags["supports_tools"] is True
# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks.
# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581).
DEEPSEEK_TEMPLATE = """
{%- if tools %}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
{%- for message in messages %}
{%- if message.role == 'assistant' and message.tool_calls %}
{%- for tc in message.tool_calls %}
{{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name +
'<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }}
{%- endfor %}
{%- endif %}
{%- endfor %}
"""
def test_detect_safetensors_features_deepseek_template_keeps_tools_on():
"""DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1")
flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE)
assert flags["supports_tools"] is True
# GLM 4.5 / 4.6 / 4.7 emit ``<tool_call>NAME\n<arg_key>...<arg_value>...
GLM_TEMPLATE = """
{%- if tools %}
For each function call, output the function name and arguments within
the following XML format:
<tool_call>{function-name}
<arg_key>{arg-key}</arg_key>
<arg_value>{arg-value}</arg_value>
</tool_call>
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
"""
def test_detect_safetensors_features_glm_template_keeps_tools_on():
"""GLM 4.x emits ``<tool_call>NAME\\n<arg_key>...``; parser handles it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6")
flags = _detect_safetensors_features(backend, GLM_TEMPLATE)
assert flags["supports_tools"] is True
# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks
# with ``functions.NAME:IDX`` as the per-call id.
KIMI_TEMPLATE = """
{%- if tools %}
<|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|>
{%- endif %}
{%- for message in messages %}
{%- if message.role == 'assistant' and message.tool_calls %}
<|tool_calls_section_begin|>
{%- for tc in message.tool_calls %}
<|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|>
{%- endfor %}
<|tool_calls_section_end|>
{%- endif %}
{%- endfor %}
"""
def test_detect_safetensors_features_kimi_template_keeps_tools_on():
"""Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct")
flags = _detect_safetensors_features(backend, KIMI_TEMPLATE)
assert flags["supports_tools"] is True
LLAMA3_2_BARE_JSON_TEMPLATE = """
{%- if tools %}
{{- 'Given the following functions, respond with JSON for a function call.' }}
{{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
{%- for message in messages %}
{%- if 'tool_calls' in message %}
{{- '{"name": "' + message.tool_calls[0].function.name + '", '}}
{{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }}
{%- endif %}
{%- endfor %}
"""
def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on():
"""Llama-3.2 bare JSON is supported, so the pill stays enabled."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE)
assert flags["supports_tools"] is True
MINICPM5_ATTRIBUTE_TEMPLATE = """
{%- if tools %}
{{- 'Available tools. Emit calls as ' }}
{{- '<function name="NAME"><parameter name="key">value</parameter></function>' }}
{%- for tool in tools %}
{{- tool | tojson }}
{%- endfor %}
{%- endif %}
"""
def test_detect_safetensors_features_attribute_function_form_keeps_tools_on():
"""The attribute form ``<function name="...">`` must be whitelisted or the pill is wrongly suppressed."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5")
flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE)
assert flags["supports_tools"] is True
def test_detect_safetensors_features_unknown_format_suppresses_tools():
"""Tools advertised with no known marker must be suppressed."""
from routes.inference import _detect_safetensors_features
tpl = (
"{%- if tools %}<|im_start|>system\n"
"Emit tool calls as JSON-RPC notifications inside the response."
"<|im_end|>{%- endif %}"
)
backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format")
flags = _detect_safetensors_features(backend, tpl)
assert flags["supports_tools"] is False
def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on():
"""Sanity check: gate only suppresses non-Qwen formats."""
"""Sanity check: Qwen <tool_call> marker still flips supports_tools."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B")
@ -454,3 +654,184 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors():
assert flags["supports_tools"] is True
assert flags["supports_reasoning"] is True
assert flags["supports_preserve_thinking"] is True
@pytest.mark.parametrize(
"opener",
[
"<｜tool▁calls▁begin｜>", # canonical
"<｜tool_calls_begin｜>", # ASCII underscores
"<｜tool▁calls｜>", # short form
"<｜tool calls begin｜>", # spaces
"<｜tool\\_calls\\_begin｜>", # escaped underscores
],
)
def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener):
# Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives
# its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ...
from routes.inference import _detect_safetensors_features
tpl = (
"{%- if tools %}tools{%- endif %}"
+ opener
+ "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}"
"<｜tool▁call▁end｜><｜tool▁calls▁end｜>"
)
backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1")
flags = _detect_safetensors_features(backend, tpl)
assert flags["supports_tools"] is True
# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON
# call form, but whose ``{"name":`` example is pretty-printed or JSON-escaped.
_WHITESPACE_BARE_JSON_TEMPLATE = (
"{%- if tools %}\n"
"To call a tool, output JSON of the form:\n"
'{ "name" : "function_name", "parameters": { } }\n'
"{%- endif %}\n"
"{{ messages }}"
)
_ESCAPED_BARE_JSON_TEMPLATE = (
"{%- if tools %}\n"
'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n'
"{%- endif %}\n"
"{{ messages }}"
)
_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = (
"{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}"
)
def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json():
# A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser
# accepts that whitespace via raw_decode.
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE)
assert flags["supports_tools"] is True
def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json():
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE)
assert flags["supports_tools"] is True
def test_detect_safetensors_features_drops_tools_when_no_parseable_form():
# Negative control: tools advertised but no parser-recognised emission form at
# all -> the pill is still dropped (the gate is not now matching everything).
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM)
assert flags["supports_tools"] is False
def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json():
# A template documenting the parser-supported {"function":...} bare-JSON alias
# must keep supports_tools, mirroring the {"name":...} form.
from routes.inference import _detect_safetensors_features
tpl = (
"{%- if tools %}\n"
'Respond with {"function": "fn", "parameters": {}}\n'
"{%- endif %}\n"
"{{ messages }}"
)
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
flags = _detect_safetensors_features(backend, tpl)
assert flags["supports_tools"] is True
# _sf_reasoning_prefill_mode gates the prefilled-<think> extractor (GGUF reasoning parity).
class TestSafetensorsReasoningPrefillGate:
# A minimal Qwen3-style template with the standard <think>/</think> markers.
_QWEN_TPL = "{% if enable_thinking %}<think>{% endif %}...</think>..."
# gemma-style bespoke reasoning channel -- no standard markers.
_GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought<channel|>"
# always-on template whose GENERATION PROMPT opens an unclosed <think> (DeepSeek-R1 / QwQ /
# Qwen3-Thinking shape): the model emits only the closing </think>, so prefill.
_ALWAYS_ON_OPEN_TPL = (
"{% for m in messages %}{{ m['content'] }}{% endfor %}"
"{% if add_generation_prompt %}<|assistant|><think>\n{% endif %}"
)
# always-on template that renders PAST assistant <think>...</think> history but leaves the
# generation prompt open with no <think> (Kimi-K2-Thinking shape): the model self-emits its
# own block, so prefill mode would blank a normal answer.
_ALWAYS_ON_HISTORY_TPL = (
"{% for m in messages %}"
"{% if m['role'] == 'assistant' %}<think>{{ m.get('reasoning_content', '') }}</think>"
"{{ m['content'] }}{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}"
)
def _features(self, **over):
base = {
"supports_reasoning": True,
"reasoning_always_on": False,
"reasoning_style": "enable_thinking",
}
base.update(over)
return base
def test_g1_enable_thinking_true(self):
# G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled.
from routes.inference import _sf_reasoning_prefill_mode
assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True
def test_g2_enable_thinking_none_defaults_on(self):
# G2: default request (None) -> prefilled (Qwen3/GLM templates default on).
from routes.inference import _sf_reasoning_prefill_mode
assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True
def test_g3_enable_thinking_false(self):
# G3: thinking explicitly off -> not prefilled.
from routes.inference import _sf_reasoning_prefill_mode
assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False
def test_g4_gpt_oss_reasoning_effort_excluded(self):
# G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(reasoning_style = "reasoning_effort")
assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
def test_g5_enable_thinking_effort_included(self):
# G5: GLM-style enable_thinking_effort also prefills.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(reasoning_style = "enable_thinking_effort")
assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True
def test_g6_non_reasoning_model(self):
# G6: no reasoning capability -> never prefilled.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(supports_reasoning = False, reasoning_style = None)
assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
def test_g7_reasoning_always_on_prompt_opens_think(self):
# G7: always-on template whose generation prompt opens <think> -> prefilled regardless of the flag.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(reasoning_always_on = True)
assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True
def test_g7b_reasoning_always_on_history_only_not_prefilled(self):
# G7b (#5704): always-on classification from rendered assistant HISTORY <think></think>
# (Kimi-K2-Thinking) whose generation prompt opens no <think>. Prefill mode would capture a
# normal answer entirely as reasoning_content and blank the visible answer, so it must be off.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(reasoning_always_on = True)
assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False
def test_g8_gemma_bespoke_channel_excluded(self):
# G8: gemma's <|think|>/<|channel> format has no </think> -> NOT prefilled
# (would otherwise swallow the whole answer as reasoning). Regression guard.
from routes.inference import _sf_reasoning_prefill_mode
assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False
def test_g9_missing_template_not_prefilled(self):
# G9: no template available -> conservative (not prefilled).
from routes.inference import _sf_reasoning_prefill_mode
assert _sf_reasoning_prefill_mode(self._features(), True, None) is False

217
studio/backend/tests/test_safetensors_reasoning_stream.py Normal file
View file

@ -0,0 +1,217 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Safetensors/MLX reasoning-block parity with GGUF.
enable_thinking templates (Qwen3/GLM) prefill an unclosed ``<think>`` so the model
emits only the closing ``</think>`` then the answer; the safetensors stream must
split the leading text into ``reasoning_content`` deltas (plain stream and tool
loop), resetting per turn and appending only visible text to the monitor. Replays a
copy of ``sf_tool_stream``'s reasoning loop against synthetic events.
"""
from __future__ import annotations
import sys
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)
from routes.inference import (
_ResponsesReasoningExtractor,
_sf_reasoning_prefill_mode,
_strip_tool_xml_for_display,
)
_THINK_TPL = "...<think>...</think>..."
_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True}
_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True}
def test_prefill_mode_on_for_enable_thinking_default():
assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True
def test_prefill_mode_off_when_thinking_disabled():
assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False
def test_prefill_mode_off_for_reasoning_effort_none():
# enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode
# would capture the whole answer as reasoning_content.
assert (
_sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none")
is False
)
assert (
_sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high")
is True
)
def test_prefill_mode_off_without_think_markers():
assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False
def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict:
"""Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content``
snapshot, feed the delta through the extractor, and reset (flushing first) on
``tool_start`` / empty ``status`` so each turn splits independently."""
prev_text = ""
extractor = _ResponsesReasoningExtractor(
parse_think_markers = True, reasoning_prefilled = prefilled
)
reasoning_deltas: list[str] = []
visible_deltas: list[str] = []
monitor: list[str] = []
tool_starts: list[dict] = []
order: list[str] = [] # sequence of ("reasoning"|"visible"|"tool_start") events
def _flush():
fr, fv = extractor.finish()
if fr:
reasoning_deltas.append(fr)
order.append("reasoning")
if fv:
visible_deltas.append(fv)
monitor.append(fv)
order.append("visible")
for event in events:
etype = event["type"]
if etype == "status":
if not event["text"]:
_flush()
prev_text = ""
extractor = _ResponsesReasoningExtractor(
parse_think_markers = True, reasoning_prefilled = prefilled
)
continue
if etype in ("tool_start", "tool_end"):
if etype == "tool_start":
_flush()
prev_text = ""
extractor = _ResponsesReasoningExtractor(
parse_think_markers = True, reasoning_prefilled = prefilled
)
tool_starts.append(event)
order.append("tool_start")
continue
clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True)
new_text = clean[len(prev_text) :]
prev_text = clean
if not new_text:
continue
r, v = extractor.feed(new_text)
if r:
reasoning_deltas.append(r)
order.append("reasoning")
if v:
visible_deltas.append(v)
monitor.append(v)
order.append("visible")
_flush()
return {
"reasoning": "".join(reasoning_deltas),
"visible": "".join(visible_deltas),
"monitor": "".join(monitor),
"tool_starts": tool_starts,
"order": order,
}
def test_s1_plain_stream_splits_prefilled_reasoning():
# S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only.
events = [
{"type": "content", "text": "Let me compute 17*23"},
{"type": "content", "text": "Let me compute 17*23 = 391</think>The answer is 391."},
]
out = _replay_sf_reasoning_stream(events, prefilled = True)
assert out["reasoning"] == "Let me compute 17*23 = 391"
assert out["visible"] == "The answer is 391."
assert out["monitor"] == "The answer is 391."
assert "<think>" not in out["reasoning"] and "</think>" not in out["visible"]
def test_s2_reasoning_flushed_before_tool_start():
# S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start.
events = [
{"type": "content", "text": "I should search"},
{"type": "content", "text": "I should search Sydney weather</think>"},
{"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"},
{"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"},
{"type": "status", "text": ""},
{"type": "content", "text": "Found it</think>Sydney is 21C today."},
]
out = _replay_sf_reasoning_stream(events, prefilled = True)
# Both turns' reasoning surfaced, answer only from turn 2.
assert "I should search Sydney weather" in out["reasoning"]
assert "Found it" in out["reasoning"]
assert out["visible"] == "Sydney is 21C today."
assert out["monitor"] == "Sydney is 21C today."
# Ordering: the pre-tool reasoning is emitted before the tool_start.
assert out["order"].index("reasoning") < out["order"].index("tool_start")
def test_s3_extractor_resets_each_turn():
# S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each).
events = [
{"type": "content", "text": "turn1 thoughts</think>partial"},
{"type": "status", "text": ""},
{"type": "content", "text": "turn2 thoughts</think>final answer"},
]
out = _replay_sf_reasoning_stream(events, prefilled = True)
assert out["reasoning"] == "turn1 thoughtsturn2 thoughts"
assert out["visible"] == "partialfinal answer"
def test_s4_harmony_full_tags_normal_mode():
# S4: gpt-oss / explicit-tag models use normal mode (prefilled=False).
events = [{"type": "content", "text": "<think>reasoning here</think>visible answer"}]
out = _replay_sf_reasoning_stream(events, prefilled = False)
assert out["reasoning"] == "reasoning here"
assert out["visible"] == "visible answer"
def test_s5_thinking_off_no_reasoning_deltas():
# S5: thinking disabled -> not prefilled, no </think>, all content is visible.
events = [{"type": "content", "text": "Just the plain answer, no thinking."}]
out = _replay_sf_reasoning_stream(events, prefilled = False)
assert out["reasoning"] == ""
assert out["visible"] == "Just the plain answer, no thinking."
assert out["monitor"] == "Just the plain answer, no thinking."
def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort():
# GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and
# enable_thinking omitted) disables thinking exactly like enable_thinking=False, so
# prefilled mode must be OFF. Otherwise the model emits no </think> and a plain
# answer is swallowed whole into reasoning_content, leaving the visible response
# empty (the exact bug: prefilled=True below eats the whole answer).
feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True}
assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False
# Thinking on (effort level or default) still prefills.
assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True
assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True
# An explicit enable_thinking=False also disables (unchanged).
assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False
# reasoning_always_on wins regardless of reasoning_effort.
always = {**feats, "reasoning_always_on": True}
assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True
# Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected.
plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True}
assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True
# End-to-end: with the corrected prefilled=False, a plain no-</think> answer is
# emitted as visible content rather than swallowed into the thinking drawer.
events = [{"type": "content", "text": "The capital of France is Paris."}]
out = _replay_sf_reasoning_stream(events, prefilled = False)
assert out["visible"] == "The capital of France is Paris."
assert out["reasoning"] == ""
# The buggy prefilled=True path is what swallowed the whole answer (guard the delta).
swallowed = _replay_sf_reasoning_stream(events, prefilled = True)
assert swallowed["visible"] == ""
assert swallowed["reasoning"] == "The capital of France is Paris."

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

File diff suppressed because it is too large Load diff

179
studio/backend/tests/test_safetensors_toolcall_wiring.py Normal file
View file

@ -0,0 +1,179 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path.
The parser and the cumulative-text state machine are already covered exhaustively by
``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the
*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and
``MLXInferenceBackend`` render the prompt through the shared
``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the
shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and
``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU
test of that seam covers the macOS MLX path too).
This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the
``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full
agentic chain end to end:
tools injected into the template -> loop parses the call -> tool dispatched once ->
tool result fed back -> generation re-entered -> final answer streamed.
It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling
end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the
tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704).
"""
from core.inference.chat_template_helpers import apply_chat_template_for_generation
from core.inference.safetensors_agentic import run_safetensors_tool_loop
TOOL_NAME = "get_weather"
TOOL_ARGS = {"city": "Paris"}
FAKE_TOOL = {
"type": "function",
"function": {
"name": TOOL_NAME,
"description": "Get the current weather for a city.",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
},
}
# Full parser matrix lives in test_safetensors_tool_loop.py.
TOOL_CALL_TEXT = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>'
FINAL_ANSWER = "The weather in Paris is sunny and 22C."
TOOL_RESULT = "Paris: sunny, 22C"
class RecordingTokenizer:
"""Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``.
Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the
real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually
forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen``
holding ``None``.
"""
def __init__(self):
self.tools_seen: list = []
self.call_count = 0
def apply_chat_template(
self,
messages,
*,
tokenize = False,
add_generation_prompt = True,
**kwargs,
):
self.call_count += 1
self.tools_seen.append(kwargs.get("tools"))
return "PROMPT"
class StubExecutor:
"""Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result.
A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run.
"""
def __init__(self, result: str):
self.result = result
self.calls: list[tuple[str, dict]] = []
def __call__(
self,
name,
arguments,
*,
cancel_event = None,
timeout = None,
session_id = None,
rag_scope = None,
disable_sandbox = False,
):
self.calls.append((name, arguments))
return self.result
def _collect(generator, max_events = 200):
events = []
for ev in generator:
events.append(ev)
if len(events) >= max_events:
break
return events
def _tool_names(tools):
return [(t.get("function") or {}).get("name") for t in (tools or [])]
def test_backend_seam_injects_tools_and_drives_full_tool_loop():
"""The shared backend seam forwards tools into the chat template, and the loop parses the call,
dispatches it once, feeds the result back, and re-enters generation for the final answer."""
tok = RecordingTokenizer()
executor = StubExecutor(TOOL_RESULT)
turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER])
active_tools_seen: list = []
conversations_seen: list = []
def single_turn(conversation, *, active_tools = None):
# Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots.
active_tools_seen.append(active_tools)
conversations_seen.append([dict(m) for m in conversation])
apply_chat_template_for_generation(tok, conversation, tools = active_tools)
text = next(turns)
mid = len(text) // 2
acc = ""
for chunk in (text[:mid], text[mid:]):
acc += chunk
yield acc
events = _collect(
run_safetensors_tool_loop(
single_turn = single_turn,
messages = [{"role": "user", "content": "What is the weather in Paris?"}],
tools = [FAKE_TOOL],
execute_tool = executor,
max_tool_iterations = 3,
)
)
# 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools).
assert tok.tools_seen, "tokenizer.apply_chat_template was never called"
assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer"
assert TOOL_NAME in _tool_names(tok.tools_seen[0])
# 2. Loop offered the tool to the first generation turn.
assert active_tools_seen and active_tools_seen[0] is not None
assert TOOL_NAME in _tool_names(active_tools_seen[0])
# 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result.
tool_starts = [e for e in events if e["type"] == "tool_start"]
tool_ends = [e for e in events if e["type"] == "tool_end"]
assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME
assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls
assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT
# 6. Final answer streams after the tool result: loop appended it and re-entered generation.
contents = [e for e in events if e["type"] == "content"]
assert contents and FINAL_ANSWER in contents[-1]["text"]
last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end")
last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content")
assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result"
# 6b. Tool result fed back into the conversation before the final turn (6 alone misses this:
# the fake generation ignores the conversation).
assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call"
final_turn_convo = conversations_seen[1]
assert any(
TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo
), "tool result was not fed back into the conversation before the final generation turn"
# 7. Guard: raw tool-call markup never leaked to the client as content.
for e in contents:
assert "<tool_call>" not in e["text"]
assert TOOL_NAME not in e["text"]

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

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

786
studio/backend/tests/test_sf_client_tools_passthrough.py Normal file
View file

@ -0,0 +1,786 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Client-tools passthrough healing for the safetensors/MLX backend.
Parity for #6801: when a NON-GGUF model is loaded and the request declares its
own ``tools`` with server-side tools OFF, text-form tool calls are promoted back
into structured ``tool_calls`` (declared tools only) via the shared healer. MLX
rides the same orchestrator path, so a single scripted backend covers both.
"""
import asyncio
import json
from types import SimpleNamespace
from models.inference import ChatCompletionRequest, ChatMessage
from routes.inference import openai_chat_completions
from core.inference.api_monitor import ApiMonitor
LOOKUP_TOOL = {
"type": "function",
"function": {
"name": "lookup",
"description": "Look something up",
"parameters": {
"type": "object",
"properties": {"q": {"type": "string"}},
"required": ["q"],
},
},
}
SEARCH_TOOL = {
"type": "function",
"function": {
"name": "search",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
}
_CALL_XML = '<tool_call>{"name": "lookup", "arguments": {"q": "cats"}}</tool_call>'
_SEARCH_XML = '<tool_call>{"name": "search", "arguments": {"query": "dogs"}}</tool_call>'
class _Request:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/chat/completions")
method = "POST"
scope: dict = {}
async def is_disconnected(self):
return False
class _ScriptedBackend:
"""Non-GGUF backend: ``generate_chat_response`` replays scripted
CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot
list for one generation, so nudge tests can vary output across turns."""
active_model_name = "sf-model"
def __init__(
self,
responder,
*,
stats = None,
):
self.models = {
"sf-model": {
"chat_template_info": {"template": "<tool_call> chatml"},
"context_length": 2048,
}
}
self._responder = responder
self._stats = stats
self.calls: list = []
self.reset_count = 0
def generate_chat_response(
self,
*,
messages,
tools = None,
stats_holder = None,
**kwargs,
):
self.calls.append({"messages": messages, "tools": tools, **kwargs})
snapshots = self._responder(messages, tools)
if stats_holder is not None and self._stats is not None:
stats_holder["stats"] = self._stats
for snap in snapshots:
yield snap
def reset_generation_state(self):
self.reset_count += 1
def _fixed(*snapshots):
"""Responder that always replays the given cumulative snapshots."""
return lambda messages, tools: list(snapshots)
def _llama_stub():
return SimpleNamespace(
is_loaded = False,
supports_tools = False,
is_vision = False,
context_length = None,
)
def _install(
monkeypatch,
backend,
*,
supports_tools = True,
):
import routes.inference as inf
from state.tool_policy import reset_tool_policy
reset_tool_policy()
monitor = ApiMonitor(max_entries = 8)
monkeypatch.setattr(inf, "api_monitor", monitor)
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub())
monkeypatch.setattr(inf, "get_inference_backend", lambda: backend)
monkeypatch.setattr(
inf,
"_detect_safetensors_features",
lambda *a, **k: {"supports_tools": supports_tools},
)
return monitor
def _request(**kwargs):
base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")])
base.update(kwargs)
return ChatCompletionRequest(**base)
def _call(payload, monkeypatch, backend, **install_kwargs):
_install(monkeypatch, backend, **install_kwargs)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
return asyncio.run(_run())
def _json_body(response):
return json.loads(response.body if hasattr(response, "body") else response.content)
def _collect_sse(response):
async def _run():
return [c async for c in response.body_iterator]
return asyncio.run(_run())
def _sse_objects(chunks):
out = []
for chunk in chunks:
if isinstance(chunk, bytes):
chunk = chunk.decode()
for line in str(chunk).splitlines():
if line.startswith("data: "):
data = line.removeprefix("data: ")
if data != "[DONE]":
out.append(json.loads(data))
return out
# ── Non-streaming ─────────────────────────────────────────────────
def test_xml_healed_to_tool_calls_non_streaming(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["content"] is None
calls = choice["message"]["tool_calls"]
assert len(calls) == 1
assert calls[0]["function"]["name"] == "lookup"
assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"}
# The client tools reached the generator (template injection).
assert backend.calls[0]["tools"] == [LOOKUP_TOOL]
def test_undeclared_call_stays_text(monkeypatch):
xml = '<tool_call>{"name": "other", "arguments": {}}</tool_call>'
backend = _ScriptedBackend(_fixed(xml))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
assert choice["message"]["content"] == xml
def test_opt_out_relays_verbatim(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
assert choice["message"]["content"] == _CALL_XML
def test_env_kill_switch_relays_verbatim(monkeypatch):
import core.inference.passthrough_healing as ph
monkeypatch.setattr(ph, "_HEALING_DISABLED", True)
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
assert choice["message"]["content"] == _CALL_XML
def test_no_tools_request_untouched(monkeypatch):
backend = _ScriptedBackend(_fixed("just a plain answer"))
payload = _request(stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
# No tools and no tool messages -> plain path, normal ChatCompletion.
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"]["content"] == "just a plain answer"
assert choice["message"].get("tool_calls") is None
def test_prose_around_call_retained(monkeypatch):
text = "Let me look:\n" + _CALL_XML + "\ndone"
backend = _ScriptedBackend(_fixed(text))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["content"] == "Let me look:\n\ndone"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
def test_empty_output_is_valid_stop(monkeypatch):
backend = _ScriptedBackend(_fixed(""))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"]["content"] in ("", None)
assert choice["message"].get("tool_calls") is None
def test_tool_role_follow_up_turn_preserves_history(monkeypatch):
backend = _ScriptedBackend(_fixed("The weather is sunny."))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "weather?"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": '{"q": "weather"}'},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"),
],
)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "The weather is sunny."
# The tool history reached the generator intact (role=tool + assistant.tool_calls).
sent = backend.calls[0]["messages"]
roles = [m["role"] for m in sent]
assert "tool" in roles
assistant = next(m for m in sent if m["role"] == "assistant")
assert assistant.get("tool_calls")
def test_dict_arguments_history_does_not_crash(monkeypatch):
# Non-spec client: assistant tool_calls[].function.arguments as a dict.
backend = _ScriptedBackend(_fixed("ok"))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "hi"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": {"q": "x"}},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"),
],
)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "ok"
def test_forced_tool_choice_narrows_promotion(monkeypatch):
# tool_choice forces `search`; a `lookup` text call must NOT promote.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(
tools = [LOOKUP_TOOL, SEARCH_TOOL],
stream = False,
tool_choice = {"type": "function", "function": {"name": "search"}},
)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
def test_parallel_cap_non_streaming(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML))
payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False)
body = _json_body(_call(payload, monkeypatch, backend))
calls = body["choices"][0]["message"]["tool_calls"]
assert len(calls) == 1
assert calls[0]["function"]["name"] == "lookup"
def test_usage_recorded_when_stats_present(monkeypatch):
stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}}
backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats)
payload = _request(tools = [LOOKUP_TOOL], stream = False)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
asyncio.run(_run())
[entry] = monitor.snapshot()
assert entry["prompt_tokens"] == 7
assert entry["completion_tokens"] == 3
# ── Nudge ─────────────────────────────────────────────────────────
def test_nudge_default_off_single_generation(monkeypatch):
# Signal present but unparseable; without opt-in, no retry.
truncated = '<tool_call>{"name": "lookup"'
backend = _ScriptedBackend(_fixed(truncated))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
_call(payload, monkeypatch, backend)
assert len(backend.calls) == 1
def test_nudge_opt_in_retry_recovers(monkeypatch):
truncated = '<tool_call>{"name": "lookup"'
def responder(messages, tools):
nudged = any(
"native tool-call format" in (m.get("content") or "")
for m in messages
if m.get("role") == "user"
)
return [_CALL_XML] if nudged else [truncated]
backend = _ScriptedBackend(responder)
payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True)
body = _json_body(_call(payload, monkeypatch, backend))
assert len(backend.calls) == 2
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
def test_nudge_double_failure_relays_original(monkeypatch):
truncated = '<tool_call>{"name": "lookup"'
backend = _ScriptedBackend(_fixed(truncated))
payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True)
body = _json_body(_call(payload, monkeypatch, backend))
assert len(backend.calls) == 2 # exactly one retry
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"]["content"] == truncated
# ── Streaming ─────────────────────────────────────────────────────
def test_streaming_heals_split_call_into_one_delta(monkeypatch):
# Cumulative snapshots that build the call across many increments.
pieces = ["<tool", '<tool_call>{"name": "loo', '<tool_call>{"name": "lookup", "argum']
cumulative = pieces + [_CALL_XML]
backend = _ScriptedBackend(_fixed(*cumulative))
payload = _request(tools = [LOOKUP_TOOL], stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert len(tool_deltas) == 1
assert tool_deltas[0]["function"]["name"] == "lookup"
finishes = [
o["choices"][0]["finish_reason"]
for o in objs
if o["choices"] and o["choices"][0].get("finish_reason")
]
assert finishes == ["tool_calls"]
def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch):
# A stream cancelled via the registry ("Stop") must NOT promote the
# buffered-but-unclosed tool markup at finalize, else it executes a tool
# the user just cancelled. Guarded on cancel_event at the finalize step.
import routes.inference as inf
cancel_id = "cancel-me-6870"
# Balanced JSON but no closing </tool_call> -> healer HOLDS it until finalize.
held = '<tool_call>{"name": "lookup", "arguments": {"q": "cats"}}'
class _CancelMidStream(_ScriptedBackend):
def __init__(self):
super().__init__(_fixed(held))
def generate_chat_response(
self,
*,
messages,
tools = None,
stats_holder = None,
**kwargs,
):
self.calls.append({"messages": messages, "tools": tools, **kwargs})
yield held # healer holds the unclosed call
inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF
backend = _CancelMidStream()
payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert tool_deltas == [] # no tool promoted after cancel
finishes = [
o["choices"][0]["finish_reason"]
for o in objs
if o["choices"] and o["choices"][0].get("finish_reason")
]
assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls
def test_streaming_no_tools_verbatim(monkeypatch):
backend = _ScriptedBackend(_fixed("hello ", "hello world"))
payload = _request(stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
text = "".join(
(o["choices"][0]["delta"].get("content") or "")
for o in objs
if o["choices"] and "delta" in o["choices"][0]
)
assert text == "hello world"
finishes = [
o["choices"][0]["finish_reason"]
for o in objs
if o["choices"] and o["choices"][0].get("finish_reason")
]
assert finishes == ["stop"]
def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch):
# Repeated then shrunk cumulative snapshots must not double-heal.
backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert len(tool_deltas) == 1
def test_streaming_parallel_cap(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML))
payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert len(tool_deltas) == 1
assert tool_deltas[0]["function"]["name"] == "lookup"
def test_streaming_generator_error_closes_cleanly(monkeypatch):
def responder(messages, tools):
raise RuntimeError("boom /secret/path")
backend = _ScriptedBackend(responder)
payload = _request(tools = [LOOKUP_TOOL], stream = True)
response = _call(payload, monkeypatch, backend)
chunks = _collect_sse(response)
joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks)
assert "An internal error occurred" in joined
assert "secret/path" not in joined # CWE-209: no path leak
assert backend.reset_count >= 1
def test_streaming_disconnect_resets_once(monkeypatch):
class _DisconnectRequest(_Request):
async def is_disconnected(self):
return True
backend = _ScriptedBackend(_fixed("a", "ab", "abc"))
payload = _request(tools = [LOOKUP_TOOL], stream = True)
_install(monkeypatch, backend)
async def _run():
resp = await openai_chat_completions(
payload, request = _DisconnectRequest(), current_subject = "u"
)
return [c async for c in resp.body_iterator]
asyncio.run(_run())
assert backend.reset_count == 1
def test_mlx_uses_same_path(monkeypatch):
# MLX and safetensors share get_inference_backend(); one scripted backend covers both.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["finish_reason"] == "tool_calls"
def test_tool_choice_none_does_not_advertise_tools(monkeypatch):
# tool_choice="none": no tools rendered into the template; history templating still applies.
backend = _ScriptedBackend(_fixed("plain answer"))
payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "plain answer"
assert backend.calls[0]["tools"] is None
def test_developer_message_folded_into_system_prompt(monkeypatch):
# The "developer" role folds into one leading system message (local templates reject it).
backend = _ScriptedBackend(_fixed("ok"))
payload = _request(
messages = [
ChatMessage(role = "developer", content = "always be terse"),
ChatMessage(role = "user", content = "hi"),
],
tools = [LOOKUP_TOOL],
stream = False,
)
_call(payload, monkeypatch, backend)
sent = backend.calls[0]["messages"]
assert sent[0]["role"] == "system"
assert "always be terse" in sent[0]["content"]
assert all(m.get("role") != "developer" for m in sent)
def test_failed_nudge_retry_keeps_original_response(monkeypatch):
# A raising retry must not 500; the first response is returned.
state = {"n": 0}
def responder(messages, tools):
state["n"] += 1
if state["n"] == 1:
return ['<tool_call>{"name":"lookup"'] # unhealable signal
raise RuntimeError("retry blew up")
backend = _ScriptedBackend(responder)
payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
assert state["n"] == 2
assert body["choices"][0]["finish_reason"] == "stop"
assert body["choices"][0]["message"]["content"] == '<tool_call>{"name":"lookup"'
def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch):
# Double-failure nudge: the first response is delivered, but the retry's
# generate() overwrites stats_holder. The monitor must record the FIRST
# attempt's usage, not the discarded retry's.
first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}}
retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}}
class _PerCallStatsBackend(_ScriptedBackend):
def __init__(self):
# Unhealable truncated markup on both attempts -> retry is discarded.
super().__init__(lambda m, t: ['<tool_call>{"name":"lookup"'])
self._stats_seq = [first_stats, retry_stats]
def generate_chat_response(
self,
*,
messages,
tools = None,
stats_holder = None,
**kwargs,
):
self.calls.append({"messages": messages, "tools": tools, **kwargs})
stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)]
if stats_holder is not None:
stats_holder["stats"] = stats
for snap in self._responder(messages, tools):
yield snap
backend = _PerCallStatsBackend()
payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
asyncio.run(_run())
assert len(backend.calls) == 2 # first attempt + one discarded retry
[entry] = monitor.snapshot()
# The delivered response is the first attempt, so its usage must be reported.
assert entry["prompt_tokens"] == 7
assert entry["completion_tokens"] == 3
def test_monitor_records_healed_call_not_raw_xml(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
asyncio.run(_run())
snap = monitor.snapshot(include_details = True)
replies = json.dumps(snap)
assert "<tool_call>" not in replies
assert "lookup" in replies
def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch):
# Monitor mirrors what the client received, never the healed-away raw markup.
backend = _ScriptedBackend(
_fixed("Sure. ", 'Sure. <tool_call>{"name": "loo', "Sure. " + _CALL_XML)
)
payload = _request(tools = [LOOKUP_TOOL], stream = True)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
response = asyncio.run(_run())
_collect_sse(response)
replies = json.dumps(monitor.snapshot(include_details = True))
assert "<tool_call>" not in replies
assert "Sure. " in replies
assert "[tool_calls] lookup(" in replies
def test_forced_tool_choice_narrows_templated_tools(monkeypatch):
# A forced function is the only schema rendered into the template.
backend = _ScriptedBackend(_fixed(_SEARCH_XML))
payload = _request(
tools = [LOOKUP_TOOL, SEARCH_TOOL],
stream = False,
tool_choice = {"type": "function", "function": {"name": "search"}},
)
body = _json_body(_call(payload, monkeypatch, backend))
templated = backend.calls[0]["tools"]
assert [t["function"]["name"] for t in templated] == ["search"]
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "search"
def test_multimodal_content_parts_flattened_for_local_template(monkeypatch):
# Remote image URLs leave image=None, so content arrives as a part LIST:
# text parts are kept, the image part dropped.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(
messages = [
ChatMessage(
role = "user",
content = [
{"type": "text", "text": "what is this?"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/cat.png"},
},
],
)
],
tools = [LOOKUP_TOOL],
stream = False,
)
body = _json_body(_call(payload, monkeypatch, backend))
templated = backend.calls[0]["messages"]
assert all(isinstance(m.get("content"), str) for m in templated)
assert any(m["content"] == "what is this?" for m in templated)
assert body["choices"][0]["finish_reason"] == "tool_calls"
def test_string_arguments_history_deserialized_for_template(monkeypatch):
# JSON-string tool_calls arguments become dicts in the templated copy;
# the HTTP response stays OpenAI-shaped.
backend = _ScriptedBackend(_fixed("done"))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "weather?"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": '{"q": "weather"}'},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"),
],
)
_json_body(_call(payload, monkeypatch, backend))
assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant")
assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"}
def test_unparseable_arguments_string_left_untouched(monkeypatch):
backend = _ScriptedBackend(_fixed("ok"))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "hi"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": "not json {"},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"),
],
)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "ok"
assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant")
assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {"
def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch):
# mcp_enabled=true with an empty registry must not silently drop the
# declared tools; the gate keys on the server-side path claiming the request.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
assert backend.calls[0]["tools"] == [LOOKUP_TOOL]

1636
studio/backend/tests/test_tool_call_parser_strict.py
View file

File diff suppressed because it is too large Load diff

76
studio/backend/tests/test_tool_strip_guard.py Normal file
View file

@ -0,0 +1,76 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""strip_tool_patterns must match the plain per-pattern loop while skipping the
quadratic no-match rescan of a closed-pair sweep whose close token is absent."""
import random
import sys
import time
from pathlib import Path
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
from core.tool_healing import (
_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS,
strip_tool_call_markup,
strip_tool_patterns,
)
def _naive(text, patterns):
for pat in patterns:
text = pat.sub("", text)
return text
_TOKENS = [
"<tool_call>",
"</tool_call>",
"<|tool_call>",
"<tool_call|>",
"<function=x>",
"<function=mcp__s__a-b>",
"</function>",
"<parameter=p>",
"</parameter>",
"call:fn{",
"}",
"{",
'<|"|>',
"A",
" ",
"\n",
"id",
"x:1",
"</tool",
"call>",
]
def test_guard_matches_plain_loop_on_fuzz():
rng = random.Random(1234)
for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS):
for _ in range(20000):
s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10)))
assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns)
def test_strip_markup_representative_cases_unchanged():
assert strip_tool_call_markup("a <tool_call>{}</tool_call> b") == "a b"
assert strip_tool_call_markup("a <function=x><parameter=p>1</parameter></function> b") == "a b"
# Non-final keeps an unclosed block; final strips it to EOF.
assert strip_tool_call_markup("a <tool_call>{partial") == "a <tool_call>{partial"
assert strip_tool_call_markup("a <tool_call>{partial", final = True) == "a"
def test_no_quadratic_blowup_on_unclosed_markers():
# Unguarded, this took minutes.
big = "<tool_call>" * 20000 + "<function=x>" * 20000
t0 = time.perf_counter()
out = strip_tool_call_markup(big, final = True)
assert time.perf_counter() - t0 < 2.0
assert out == ""

643
studio/backend/tests/test_tool_xml_strip.py
View file

@ -24,17 +24,74 @@ import re as _re
_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text()
_m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL)
assert _m, "could not extract _TOOL_XML_RE source"
_ns = {"_re": _re}
# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped;
# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here.
assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group(
1
), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)"
# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted
# ``_re.compile`` expression resolves the same source.
from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
from core.inference.tool_call_parser import (
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
_strip_glm_calls,
_strip_mistral_closed_calls,
)
from typing import Optional as _Optional
_ns = {
"_re": _re,
"_DS_OPEN_SRC": _DS_OPEN_SRC,
"Optional": _Optional,
"_strip_mistral_closed_calls": _strip_mistral_closed_calls,
"_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls,
"_strip_glm_calls": _strip_glm_calls,
"_strip_function_xml_calls": _strip_function_xml_calls,
}
exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns)
_TOOL_XML_RE = _ns["_TOOL_XML_RE"]
_helper = _re.search(
r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n"
r"(?: .+\n)+",
# The display helper uses the closed-only variant before the last think block; keep it in scope.
_mc = _re.search(r"_TOOL_XML_CLOSED_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL)
assert _mc, "could not extract _TOOL_XML_CLOSED_RE source"
exec(f"_TOOL_XML_CLOSED_RE = _re.compile({_mc.group(1)})", _ns)
_TOOL_XML_CLOSED_RE = _ns["_TOOL_XML_CLOSED_RE"]
# Signatures may span multiple lines and now carry the enabled_tool_names gate; match
# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body.
_xml_helper = _re.search(
r"def _strip_tool_xml\((?:.|\n)*?\) -> str:\n(?: .+\n)+",
_src,
)
assert _helper, "could not extract _strip_tool_xml_for_display source"
assert _xml_helper, "could not extract _strip_tool_xml source"
assert "_strip_mistral_closed_calls" in _xml_helper.group(
0
), "extracted _strip_tool_xml no longer runs the Mistral balanced strip"
exec(_xml_helper.group(0), _ns)
_strip_tool_xml = _ns["_strip_tool_xml"]
# Extract the gate helper and display strip up to the next top-level ``logger =``.
_helper = _re.search(
r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)",
_src,
_re.DOTALL,
)
assert _helper, "could not extract display strip helper source"
# The extracted block spans _display_tool_name_gate through _strip_tool_xml (defined before
# ``logger =``); confirm the shared _strip_tool_xml delegate is present.
assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates"
exec(_helper.group(0), _ns)
_strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"]
_display_tool_name_gate = _ns["_display_tool_name_gate"]
_gate_src = _re.search(
r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+",
_src,
)
assert _gate_src, "could not extract _gemma_strip_gate source"
exec(_gate_src.group(0), _ns)
_gemma_strip_gate = _ns["_gemma_strip_gate"]
# ── Well-formed pairs ─────────────────────────────────────────────
@ -46,6 +103,66 @@ def test_route_display_strip_respects_disabled_auto_heal_contract():
assert "<tool_call>" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
def test_route_display_strip_preserves_rehearsal_inside_think():
# A rehearsed bracket call inside think is reasoning: the block is preserved while a real
# call outside it still strips.
text = '<think>plan: search[ARGS]{"q":"x"}</think> answer [TOOL_CALLS]web_search{"q":"y"} tail'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert '<think>plan: search[ARGS]{"q":"x"}</think>' in out
assert "[TOOL_CALLS]web_search" not in out
assert "answer" in out and "tail" in out
def test_route_display_strip_keeps_bare_args_before_think_block():
# A bare ``foo[ARGS]`` before a think block is prose: EOS-anchored tail arms run only on
# the last segment (earlier segments use the closed-only regex).
text = "Please pass foo[ARGS] <think>pause</think> to the template."
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) == text
def test_route_display_strip_removes_complete_call_before_think_block():
# A complete bracket call before a think block still strips (balanced scan runs on every segment).
text = 'before search[ARGS]{"q":"x"} <think>pause</think> after'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "search[ARGS]" not in out
assert "<think>pause</think>" in out
assert "before" in out and "after" in out
def test_route_display_strip_removes_closed_xml_before_think_block():
# A closed <tool_call> before a think block is removed in the non-last segment.
text = 'pre <tool_call>{"name":"x"}</tool_call> <think>p</think> tail'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "<tool_call>" not in out
assert "<think>p</think>" in out
assert "pre" in out and "tail" in out
def test_all_route_cleanup_sites_use_protected_display_helper():
# Every route cleanup site must use _strip_tool_xml_for_display (think-preserving,
# balanced); raw _TOOL_XML_RE.sub corrupted think rehearsal and trailing prose. The only
# legitimate raw sub lives inside the helper itself.
raw_sub_lines = [
(i, line)
for i, line in enumerate(_src.splitlines(), 1)
if "_TOOL_XML_RE.sub(" in line and not line.lstrip().startswith("#")
]
assert len(raw_sub_lines) == 1, (
"raw _TOOL_XML_RE.sub must appear only inside _strip_tool_xml_for_display; "
f"found extra call sites: {raw_sub_lines!r}"
)
def test_route_display_strip_removes_mistral_tool_calls_with_nested_json():
# _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral
# balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON).
text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail'
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[TOOL_CALLS]" not in out and "web_search" not in out, out
assert out == "ok tail"
def test_strips_well_formed_tool_call():
text = (
"Let me search.\n"
@ -73,6 +190,26 @@ def test_strips_function_only_well_formed():
assert "Done." in cleaned
def test_strips_function_attribute_form():
# Attribute form ``<function name="...">`` (MiniCPM-5 / MiniMax-M2) must strip from the route too
# (it previously leaked into the UI); a dotted/hyphenated name also strips.
text = (
'Sure.\n<function name="get_weather">\n'
"<parameter=city>\nSydney\n</parameter>\n</function>\nDone."
)
cleaned = _TOOL_XML_RE.sub("", text)
assert "<function name=" not in cleaned
assert "</function>" not in cleaned
assert "Sure." in cleaned and "Done." in cleaned
dotted = 'A <function name="srv.list-issues">x</function> B'
assert _TOOL_XML_RE.sub("", dotted) == "A B"
# Auto-Heal-disabled display contract still preserves literal markup.
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
assert "<function name=" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
# ── Orphan openings ───────────────────────────────────────────────
@ -155,6 +292,32 @@ def test_strips_tail_only_parameter_orphan_no_trailing_ws():
assert "Final answer." in cleaned
def test_strips_complete_bracket_tag_keeps_trailing_prose():
# A complete Mistral call strips only its balanced JSON, leaving following prose intact.
cleaned = _TOOL_XML_RE.sub("", '[TOOL_CALLS]web_search{"q":"x"} and then prose')
assert "[TOOL_CALLS]" not in cleaned
assert "and then prose" in cleaned
def test_strips_unclosed_bracket_tail():
# Close brace lost to EOS: the truncated tail strips to the end instead of leaking.
cleaned = _TOOL_XML_RE.sub("", 'here [TOOL_CALLS]web_search{"query":"weather"')
assert "[TOOL_CALLS]" not in cleaned
assert cleaned.strip() == "here"
def test_strips_unclosed_rehearsal_tail():
cleaned = _TOOL_XML_RE.sub("", 'text python[ARGS]{"code":"print(1)"')
assert "[ARGS]" not in cleaned
assert cleaned.strip() == "text"
def test_strips_hyphenated_mcp_bracket_name():
cleaned = _TOOL_XML_RE.sub("", 'x [TOOL_CALLS]mcp__srv__list-issues{"q":"x"}')
assert "list-issues" not in cleaned
assert cleaned.strip() == "x"
def test_preserves_mid_string_parameter_in_code_sample():
# Tail-anchor on `</parameter>` so doc/example prose survives.
text = (
@ -281,3 +444,473 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam():
elapsed = time.perf_counter() - t0
assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens"
assert "<tool_call>" not in cleaned
# ── Two-level-nested bracket JSON (balanced-scan strip) ──────────
def test_route_strip_two_level_nested_bracket_keeps_trailing_prose():
# Two-level-nested args must be removed whole so the trailing prose survives.
text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after'
cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert cleaned == "before after"
assert "[TOOL_CALLS]" not in cleaned
def test_route_strip_two_level_nested_rehearsal_keeps_trailing_prose():
text = 'note python[ARGS]{"a":{"b":{"c":1}}} done'
cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert cleaned == "note done"
assert "[ARGS]" not in cleaned
def test_route_strip_removes_call_with_literal_think_in_argument():
# A literal <think> inside a call argument strips with the call, not as reasoning.
text = (
'<tool_call>{"name":"write","arguments":'
'{"text":"compare <think> and </think> tags"}}</tool_call>'
)
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "<tool_call>" not in out and '"name"' not in out
def test_route_strip_removes_truncated_mistral_array():
# A canonical array truncated by EOS is stripped by the route fallback like other orphans.
text = 'before [TOOL_CALLS] [{"name":"a","arguments":{"x":1}}' # missing ]
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[TOOL_CALLS]" not in out and "{" not in out
assert "before" in out
def test_route_strip_keeps_prose_mentioning_args_marker():
# ``foo[ARGS] in a sentence`` is prose; the rehearsal arm must not truncate the line.
text = "Please pass foo[ARGS] to the template and continue reading."
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert out == text
def test_route_strip_handles_mistral_v11_call_id_args_shape():
# v11 [CALL_ID]/[ARGS] shape (Mistral Small 3.2) must strip whole.
text = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out
assert "before" in out and "after" in out
# ── Mistral [/TOOL_CALLS] closer + literal <think> inside a call ───────────────
from core.tool_healing import strip_tool_call_markup as _strip_tool_call_markup
def test_core_strip_removes_orphan_tool_calls_closer_array_form():
# The bare v11 [/TOOL_CALLS] closer left by the balanced scan must not leak as content.
text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]'
assert _strip_tool_call_markup(text, final = True) == ""
def test_core_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail():
text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail'
assert _strip_tool_call_markup(text, final = True) == "tail"
def test_core_strip_removes_call_with_literal_think_in_argument():
# An unclosed literal <think> inside call arguments strips with the call (argument data).
text = 'before <tool_call>{"name":"write","arguments":{"text":"literal <think> marker"}}</tool_call> after'
assert _strip_tool_call_markup(text, final = True) == "before after"
def test_route_display_strip_removes_orphan_tool_calls_closer_array_form():
text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert out.strip() == ""
def test_route_display_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail():
text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[/TOOL_CALLS]" not in out
assert out.strip() == "tail"
def test_incomplete_xml_call_with_literal_think_in_arg_is_stripped():
# An incomplete <tool_call> holding a literal <think> strips to EOS, not as a reasoning
# block (the unclosed tail _tool_call_markup_spans previously missed).
from core.tool_healing import parse_tool_calls_from_text as _parse
from core.tool_healing import strip_tool_call_markup as _strip
text = 'before <tool_call>{"name":"write","arguments":{"text":"literal <think> marker"}} after'
assert [c["function"]["name"] for c in _parse(text)] == ["write"]
assert _strip(text, final = True) == "before"
# A real reasoning block with no tool call is still preserved verbatim.
assert (
_strip("answer <think>real</think> done", final = True) == "answer <think>real</think> done"
)
# A complete call followed by a real reasoning block: call stripped, block kept.
mixed = '<tool_call>{"name":"a","arguments":{}}</tool_call> mid <think>r</think> end'
assert _strip(mixed, final = True) == "mid <think>r</think> end"
# ── enabled-tool gate for the ambiguous bare-rehearsal strip (#5704) ──
def test_display_tool_name_gate_returns_active_names_or_none():
# Empty / no tools -> None (unrestricted; keep the legacy strip-all behavior).
assert _display_tool_name_gate([]) is None
assert _display_tool_name_gate(None) is None
# OpenAI-shaped tool dicts -> set of function names, malformed entries dropped.
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "run_python"}},
{"type": "function"}, # no name
{"nope": 1}, # no function
]
assert _display_tool_name_gate(tools) == {"web_search", "run_python"}
def test_route_display_strip_keeps_inactive_rehearsal_when_gated():
# P1 #5704: an inactive ``foo[ARGS]{...}`` is prose; the gated strip leaves the sentence intact.
gate = {"web_search"}
text = 'foo[ARGS]{"x":1} is just syntax.'
assert (
_strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate)
== text
)
# A bare marker with no JSON body is likewise prose when inactive.
assert (
_strip_tool_xml_for_display(
"use foo[ARGS] here", auto_heal_tool_calls = True, enabled_tool_names = gate
)
== "use foo[ARGS] here"
)
def test_route_display_strip_removes_active_rehearsal_when_gated():
# Mirror case: an active tool name is a real rehearsal and still strips.
gate = {"web_search"}
out = _strip_tool_xml_for_display(
'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate
)
assert "web_search[ARGS]" not in out
assert out.strip() == "done"
def test_route_display_strip_ungated_strips_all_rehearsal_unchanged():
# Backwards-compat: with no gate (None) the bare rehearsal strips as before.
text = 'foo[ARGS]{"x":1} is just syntax.'
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax."
assert (
_strip_tool_xml_for_display(
text, auto_heal_tool_calls = True, enabled_tool_names = None
).strip()
== "is just syntax."
)
def test_route_display_strip_control_token_stripped_regardless_of_gate():
# [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate.
gate = {"web_search"}
out = _strip_tool_xml_for_display(
'[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate
)
assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out
assert out.strip() == "keep"
def test_core_strip_gates_bare_rehearsal_on_enabled_tools():
# P1 (#5704): the shared strip gate mirrors the parse gate -- inactive names are prose
# and preserved, active names strip, ``None`` keeps legacy strip-all.
from core.tool_healing import strip_tool_call_markup as _strip
text = 'foo[ARGS]{"x":1} is just syntax.'
assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text
assert (
_strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"})
== "done"
)
assert _strip(text, final = True).strip() == "is just syntax."
assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax."
def test_route_display_strip_gate_preserves_inactive_history_rehearsal():
# The GGUF history sanitiser passes the gate, so a documented inactive shape survives in
# the replayed prompt context.
gate = _display_tool_name_gate([{"function": {"name": "web_search"}}])
text = 'To call it write foo[ARGS]{"x":1} in your reply.'
assert 'foo[ARGS]{"x":1}' in _strip_tool_xml_for_display(
text, auto_heal_tool_calls = True, enabled_tool_names = gate
)
# An ACTIVE name is still stripped as a real rehearsed call.
assert "web_search[ARGS]" not in _strip_tool_xml_for_display(
'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate
)
# No gate (legacy) strips every NAME[ARGS]{...}.
assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate():
# Wiring guard: the GGUF history strip must forward the display gate like the live strip.
block = _re.search(
r"Strip stale tool-call XML from conversation history.*?\.strip\(\)",
_src,
_re.DOTALL,
)
assert block, "could not locate GGUF history sanitizer block"
assert "enabled_tool_names" in block.group(
0
), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display"
def test_route_history_and_passthrough_forward_the_display_gate():
# The safetensors/Anthropic history sanitisers and the Anthropic non-stream passthrough
# must forward the gate so inactive examples survive in replayed prompt / final text.
blocks = {
"safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)",
"anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)",
"anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)",
}
for label, pat in blocks.items():
m = _re.search(pat, _src, _re.DOTALL)
assert m, f"could not locate {label} strip block"
assert "enabled_tool_names" in m.group(
0
), f"{label} must forward enabled_tool_names to _strip_tool_xml_for_display"
# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ──
def test_strips_deepseek_space_opener_variant():
# The space-separated opener is parsed by the parser, so the display strip
# must remove it too (the shared opener alternation is reused here).
text = (
"pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>"
'{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post'
)
cleaned = _TOOL_XML_RE.sub("", text)
assert "tool" not in cleaned.replace("post", "").replace("pre", "")
assert cleaned == "pre post"
def test_strips_deepseek_escaped_underscore_opener_variant():
text = (
"pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>"
'{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post'
)
cleaned = _TOOL_XML_RE.sub("", text)
assert cleaned == "pre post"
def test_strips_bare_kimi_call_without_section_wrapper():
# Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no
# section wrapper; the parser accepts it, so the strip must cover it.
text = (
"pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>"
'{"a":1}<|tool_call_end|> post'
)
cleaned = _TOOL_XML_RE.sub("", text)
assert "tool_call_begin" not in cleaned
assert cleaned == "pre post"
@pytest.mark.parametrize(
"text",
[
# Prose that merely names a Kimi/DeepSeek marker (no real call follows) must
# survive: the call-shaped lookahead fires only on a real call or a bare EOF
# fragment, so an answer discussing the protocol is never truncated.
"See <|tool_call_begin|> in the docs. More prose after it.",
"The <|tool_calls_section_begin|> marker opens a batch. Read on.",
"DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.",
],
)
def test_deepseek_kimi_false_alarm_prose_is_kept(text):
# Regression for the route arm truncating a prose answer that references a marker
# without a following call (parser _TOOL_ALL_PATS already had this lookahead).
assert _TOOL_XML_RE.sub("", text) == text
def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix():
# The lookahead must not weaken real-call stripping: closed, truncated, and bare
# EOF-fragment forms all still get removed.
closed = (
"answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>"
'{"a":1}<|tool_call_end|> tail'
)
assert _TOOL_XML_RE.sub("", closed) == "answer tail"
eof_fragment = "prefix <|tool_call_begin|>"
assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix "
deepseek = (
"reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>"
'{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>'
)
assert _TOOL_XML_RE.sub("", deepseek) == "reply "
# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ──────
# Llama-3 <|python_tag|> arm bounds on REAL sentinels only
def test_python_tag_strip_consumes_literal_sentinel_in_arg():
# A <|python_tag|> tool call whose JSON argument carries a literal <|...|>
# token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped
# at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display.
text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}'
cleaned = _TOOL_XML_RE.sub("", text)
assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}"
@pytest.mark.parametrize(
"sentinel",
[
"<|eot_id|>",
"<|eom_id|>",
"<|start_header_id|>",
"<|end_header_id|>",
],
)
def test_python_tag_strip_stops_at_real_sentinel(sentinel):
# A genuine Llama control sentinel still bounds the strip so following
# assistant text is preserved (the arm must not swallow past it).
text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer'
cleaned = _TOOL_XML_RE.sub("", text)
assert (
cleaned == f"{sentinel}visible answer"
), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}"
def test_python_tag_strip_restarts_on_second_python_tag():
# A second <|python_tag|> opens a new tool-call region, so the whole pair is
# stripped (the arm bounds the first, then the next match consumes the rest).
text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}'
cleaned = _TOOL_XML_RE.sub("", text)
assert cleaned == "", f"second python_tag region leaked: {cleaned!r}"
def test_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole():
# GLM 4.x emits <tool_call>NAME<arg_key>k</arg_key><arg_value>v</arg_value> ...</tool_call>.
text = (
"<tool_call>web_search\n<arg_key>query</arg_key>\n"
"<arg_value>find </tool_call> here</arg_value>\n</tool_call> done"
)
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "</arg_value>" not in out
assert "<arg_key>" not in out
assert out.strip() == "done"
def test_glm_normal_and_qwen_calls_still_stripped_by_route():
# Regression: a normal GLM call (no literal close tag) and a Qwen
# <tool_call>{json}</tool_call> are still stripped; trailing prose is kept.
glm = "<tool_call>get_time\n<arg_key>tz</arg_key>\n<arg_value>UTC</arg_value>\n</tool_call> ok"
assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok"
qwen = '<tool_call>{"name":"web_search","arguments":{"q":"x"}}</tool_call> after'
assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after"
def test_route_strip_removes_param_alias_close_tag():
# The parser accepts the <param name="...">...</param> attribute-form alias of
# <parameter=...>; the route tail cleanup must strip an orphan </param> close too.
assert _strip_tool_xml_for_display("answer </param>", auto_heal_tool_calls = True) == "answer "
assert (
_strip_tool_xml_for_display("answer </parameter>", auto_heal_tool_calls = True) == "answer "
)
def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup():
# A literal <function=...></function> in a value must not truncate the strip: the route runs the
# parser's guarded function-XML scan before the regex, matching the core strip.
text = "<function=python><parameter=code><function=evil></function></parameter></function> tail"
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail"
def test_route_strip_gates_wrapperless_gemma_by_enabled_tools():
# The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names,
# like the parser/loop, so a disabled/example name in prose is preserved in ...
prose = "To document syntax you write call:foo{query:example}. That shows the format."
assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"})
# An enabled name is still a real call and stripped.
assert "call:web_search" not in _strip_tool_xml(
"Answer. call:web_search{query:x}", {"web_search"}
)
# No gate (legacy) strips every closed call.
assert "call:foo" not in _strip_tool_xml(prose)
def test_gemma_strip_gate_empty_tools_preserves_prose():
# With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls
# back to strip-all and deletes an answer that documents the call:NAME{...} syntax.
assert _gemma_strip_gate([]) == set()
assert _gemma_strip_gate(None) == set()
assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"}
prose = "To document syntax you write call:foo{query:example}. That shows the format."
assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([]))
assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None))
# An enabled tool's real call is still stripped.
assert "call:web_search" not in _strip_tool_xml(
"Answer. call:web_search{query:x}",
_gemma_strip_gate([{"function": {"name": "web_search"}}]),
)
def test_strip_keeps_prose_after_closed_function_call_with_literal_close():
# The call ends at its first non-data close: prose after it survives the
# strip even when it mentions a literal </function>.
from core.inference.tool_call_parser import strip_tool_markup
text = (
"<function=web_search><parameter=query>cats</parameter></function>"
" Done. The tag </function> closes a call."
)
assert strip_tool_markup(text, final = True) == "Done. The tag </function> closes a call."
def test_final_strip_keeps_prose_mentioning_bare_markers():
# A false-alarm marker in a normal answer must not lose everything after
# it; only text that looks like that family's call start drops.
from core.inference.tool_call_parser import strip_tool_markup
for text in (
"See [TOOL_CALLS] docs for details. More prose after.",
"<|python_tag|> is the Llama marker. Explanation continues.",
"The <|tool_call> opener wraps Gemma calls.",
):
assert strip_tool_markup(text, final = True) == text
# A bare marker at end-of-text is a fragment and still drops.
assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text"
def test_final_strip_still_drops_truncated_marker_calls():
from core.inference.tool_call_parser import strip_tool_markup
for text in (
'[TOOL_CALLS][{"name":"web_search","argu',
'[TOOL_CALLS]web_search[ARGS]{"q":"x',
'<|python_tag|>{"name":"web_search","par',
'<|python_tag|>foo.call(items=["a',
"<|tool_call>call:web_search{query:tru",
):
assert strip_tool_markup(text, final = True) == ""
def test_chained_bare_json_strip_consumes_all_calls():
# The loops keep this text as next-turn history: a leftover executed call
# would be replayed alongside the structured tool_calls.
from core.inference.tool_call_parser import strip_leading_bare_json_call
enabled = {"web_search", "python"}
chained = (
'{"name":"web_search","parameters":{"q":"first"}};'
'{"name":"python","parameters":{"code":"x"}}'
)
assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == ""
assert (
strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled)
== "trailing prose"
)
# The chain stops at a non-call answer object, which stays visible.
call_then_answer = (
'{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}'
)
assert (
strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled)
== '{"name":"web_search","result":"data"}'
)

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

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

21
studio/backend/utils/mlx_repair.py
View file

@ -45,9 +45,21 @@ DISABLE_ENV_VAR = "UNSLOTH_DISABLE_MLX_AUTOREPAIR"
# deps). mlx-vlm especially must be >=0.4.4: an older one still imports but
# breaks VLM Train/Export, so installing it would wrongly clear chat-only.
_MLX_MIN_VERSIONS = {"mlx": "0.22.0", "mlx-lm": "0.22.0", "mlx-vlm": "0.4.4"}
# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights
# rejects q_norm/k_norm, so a self-heal must not pull it. mlx-lm #1242.
_MLX_BAD_VERSIONS = {"mlx-lm": ("0.31.3",)}
_MLX_PACKAGE_NAMES = tuple(_MLX_MIN_VERSIONS)
_MLX_RUNTIME_IMPORTS = ("mlx.core", "mlx_lm", "mlx_lm.sample_utils", "mlx_vlm")
MLX_PACKAGES = tuple(f"{name}>={version}" for name, version in _MLX_MIN_VERSIONS.items())
def _mlx_spec(name: str, version: str) -> str:
spec = f"{name}>={version}"
for bad in _MLX_BAD_VERSIONS.get(name, ()):
spec += f",!={bad}"
return spec
MLX_PACKAGES = tuple(_mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items())
_MLX_REINSTALL_ARGS = tuple(
arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name)
)
@ -140,7 +152,12 @@ def _mlx_versions_satisfy_minimums() -> bool:
return False
for name, minimum in _MLX_MIN_VERSIONS.items():
try:
if Version(_dist_version(name)) < Version(minimum):
installed = Version(_dist_version(name))
if installed < Version(minimum):
return False
# A known-broken build counts as unsatisfied so the self-heal
# reinstalls a good one; Version compare matches 0.31.3(.0/+local).
if any(installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())):
return False
except PackageNotFoundError:
return False

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

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

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

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

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

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

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

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

10
studio/frontend/src/components/floating-monitor.tsx
View file

@ -27,9 +27,11 @@ function usageTextClass(percent: number): string {
return "text-primary";
}
function formatGb(value: number): string {
function formatGiB(value: number): string {
// RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching
// nvidia-smi and PyTorch, so label the readout GiB rather than GB.
const digits = value >= 10 ? 1 : 2;
return `${value.toFixed(digits)} GB`;
return `${value.toFixed(digits)} GiB`;
}
export function FloatingMonitor() {
@ -116,7 +118,7 @@ export function FloatingMonitor() {
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGb(ramUsed)} / {formatGb(ramTotal)}
{formatGiB(ramUsed)} / {formatGiB(ramTotal)}
</div>
<Progress
value={ramPercent}
@ -144,7 +146,7 @@ export function FloatingMonitor() {
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGb(vramUsed)} / {formatGb(vramTotal)}
{formatGiB(vramUsed)} / {formatGiB(vramTotal)}
</div>
<Progress
value={vramPercent}

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

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

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

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

1
studio/frontend/src/features/chat/api/chat-settings-api.ts
View file

@ -26,6 +26,7 @@ export interface PersistedChatSettings {
collapseHtmlArtifacts?: boolean;
allowArtifactNetworkAccess?: boolean;
autoHealToolCalls?: boolean;
nudgeToolCalls?: boolean;
maxToolCallsPerMessage?: number;
toolCallTimeout?: number;
}

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

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

25
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -1732,6 +1732,7 @@ export function ChatSettingsPanel({
<CollapsibleSection label="Tools">
<div className="flex flex-col gap-5 pt-1">
<AutoHealToolCallsToggle />
<NudgeToolCallsToggle />
<ConfirmToolCallsToggle />
<BypassPermissionsToggle />
<MaxToolCallsSlider />
@ -2006,6 +2007,30 @@ function AutoHealToolCallsToggle() {
);
}
function NudgeToolCallsToggle() {
const nudgeToolCalls = useChatRuntimeStore((s) => s.nudgeToolCalls);
const setNudgeToolCalls = useChatRuntimeStore((s) => s.setNudgeToolCalls);
return (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Nudge Tool Calls
</span>
<InfoHint>
When a tool call cannot be repaired, re-ask the model once so the
intended tool still runs. API requests stay opt-in.
</InfoHint>
</div>
<Switch
className="panel-switch"
checked={nudgeToolCalls}
onCheckedChange={setNudgeToolCalls}
/>
</div>
);
}
function ConfirmToolCallsToggle() {
const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls);
const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls);

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

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

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

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

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

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

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

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

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

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

14
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -646,6 +646,7 @@ type ChatRuntimeStore = {
toolStatus: string | null;
generatingStatus: string | null;
autoHealToolCalls: boolean;
nudgeToolCalls: boolean;
maxToolCallsPerMessage: number;
toolCallTimeout: number;
kvCacheDtype: string | null;
@ -780,6 +781,7 @@ type ChatRuntimeStore = {
setGeneratingStatus: (status: string | null) => void;
setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void;
setAutoHealToolCalls: (enabled: boolean) => void;
setNudgeToolCalls: (enabled: boolean) => void;
setMaxToolCallsPerMessage: (value: number) => void;
setToolCallTimeout: (value: number) => void;
setKvCacheDtype: (dtype: string | null) => void;
@ -832,6 +834,7 @@ type ScalarSettingKey =
| "collapseHtmlArtifacts"
| "allowArtifactNetworkAccess"
| "autoHealToolCalls"
| "nudgeToolCalls"
| "maxToolCallsPerMessage"
| "toolCallTimeout";
@ -869,6 +872,7 @@ const SCALAR_SETTING_KEYS = [
"collapseHtmlArtifacts",
"allowArtifactNetworkAccess",
"autoHealToolCalls",
"nudgeToolCalls",
"maxToolCallsPerMessage",
"toolCallTimeout",
] as const satisfies readonly ScalarSettingKey[];
@ -1103,6 +1107,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
generatingStatus: null,
activeDiffusionCanvas: null,
autoHealToolCalls: true,
nudgeToolCalls: true,
maxToolCallsPerMessage: 25,
toolCallTimeout: 5,
kvCacheDtype: null,
@ -1544,6 +1549,15 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
);
return { autoHealToolCalls };
}),
setNudgeToolCalls: (nudgeToolCalls) =>
set((state) => {
setScalarSettingVersion(
"nudgeToolCalls",
nudgeToolCalls,
state.nudgeToolCalls,
);
return { nudgeToolCalls };
}),
setMaxToolCallsPerMessage: (maxToolCallsPerMessage) =>
set((state) => {
setScalarSettingVersion(

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

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

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

@ -357,6 +357,7 @@ export interface OpenAIChatCompletionsRequest {
context_length?: number;
};
auto_heal_tool_calls?: boolean;
nudge_tool_calls?: boolean;
max_tool_calls_per_message?: number;
tool_call_timeout?: number;
session_id?: string;

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

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

10
studio/frontend/src/features/chat/utils/chat-settings-storage.ts
View file

@ -21,6 +21,7 @@ import type { ReasoningEffort } from "../stores/chat-runtime-store";
const AUTO_TITLE_KEY = "unsloth_chat_auto_title";
const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls";
const NUDGE_TOOL_CALLS_KEY = "unsloth_nudge_tool_calls";
const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message";
const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout";
const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params";
@ -223,6 +224,7 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings {
value.allowArtifactNetworkAccess,
);
const autoHealToolCalls = sanitizeBool(value.autoHealToolCalls);
const nudgeToolCalls = sanitizeBool(value.nudgeToolCalls);
const maxToolCallsPerMessage = sanitizeInt(value.maxToolCallsPerMessage, 1);
const toolCallTimeout = sanitizeInt(value.toolCallTimeout, 1);
@ -245,6 +247,9 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings {
if (autoHealToolCalls !== undefined) {
settings.autoHealToolCalls = autoHealToolCalls;
}
if (nudgeToolCalls !== undefined) {
settings.nudgeToolCalls = nudgeToolCalls;
}
if (maxToolCallsPerMessage !== undefined) {
settings.maxToolCallsPerMessage = maxToolCallsPerMessage;
}
@ -305,6 +310,7 @@ export function isEmptyChatSettings(settings: PersistedChatSettings): boolean {
settings.collapseHtmlArtifacts === undefined &&
settings.allowArtifactNetworkAccess === undefined &&
settings.autoHealToolCalls === undefined &&
settings.nudgeToolCalls === undefined &&
settings.maxToolCallsPerMessage === undefined &&
settings.toolCallTimeout === undefined
);
@ -335,6 +341,7 @@ export function loadLegacyChatSettings(): PersistedChatSettings {
const collapseHtmlArtifacts = loadBool(COLLAPSE_HTML_ARTIFACTS_KEY);
const allowArtifactNetworkAccess = loadBool(ALLOW_ARTIFACT_NETWORK_ACCESS_KEY);
const autoHealToolCalls = loadBool(AUTO_HEAL_TOOL_CALLS_KEY);
const nudgeToolCalls = loadBool(NUDGE_TOOL_CALLS_KEY);
const maxToolCallsPerMessage = loadInt(MAX_TOOL_CALLS_KEY, 1);
const toolCallTimeout = loadInt(TOOL_CALL_TIMEOUT_KEY, 1);
const allCustomPresets = sanitizeCustomPresets([
@ -361,6 +368,9 @@ export function loadLegacyChatSettings(): PersistedChatSettings {
if (autoHealToolCalls !== undefined) {
settings.autoHealToolCalls = autoHealToolCalls;
}
if (nudgeToolCalls !== undefined) {
settings.nudgeToolCalls = nudgeToolCalls;
}
if (maxToolCallsPerMessage !== undefined) {
settings.maxToolCallsPerMessage = maxToolCallsPerMessage;
}

4
studio/frontend/src/features/hub/hub-page.tsx
View file

@ -1085,11 +1085,11 @@ export function ModelsPage() {
const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu);
const gpuLabel = gpu.available
? `${Math.round(gpu.memoryTotalGb)} GB`
? `${Math.round(gpu.memoryTotalGb)} GiB`
: "Unavailable";
const ramLabel =
gpu.systemRamTotalGb > 0
? `${Math.round(gpu.systemRamTotalGb)} GB`
? `${Math.round(gpu.systemRamTotalGb)} GiB`
: "Unavailable";
const coreLabel =
gpu.cpuCore > 0 && gpu.cpuThread > 0

2
studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
View file

@ -125,7 +125,7 @@ export function SummaryStep() {
<span className="text-xs text-muted-foreground">GPU</span>
<div className="flex items-center gap-2">
<span className="text-sm font-medium">{hw.gpuName ?? "---"}</span>
<Badge variant="secondary">{hw.vramTotalGb != null ? `${hw.vramTotalGb} GB` : "---"}</Badge>
<Badge variant="secondary">{hw.vramTotalGb != null ? `${hw.vramTotalGb} GiB` : "---"}</Badge>
</div>
</div>
</div>

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

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

2
studio/frontend/src/features/settings/tabs/about-tab.tsx
View file

@ -158,7 +158,7 @@ export function AboutTab() {
<code className="font-mono text-xs text-muted-foreground">
{gpu.name ?? "—"}
{gpu.vramTotalGb != null
? ` · ${Math.round(gpu.vramTotalGb)} GB`
? ` · ${Math.round(gpu.vramTotalGb)} GiB`
: ""}
</code>
</SettingsRow>

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

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

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

@ -42,6 +42,7 @@ import {
updatePreviewSharing,
} from "../api/preview-sharing";
import {
EmbeddingModelBlockedError,
type EmbeddingModelSettings,
EmbeddingModelVerificationError,
loadEmbeddingModelSettings,
@ -78,6 +79,7 @@ const PREFS_KEYS: string[] = [
"unsloth_chat_auto_title",
"unsloth_hf_token",
"unsloth_auto_heal_tool_calls",
"unsloth_nudge_tool_calls",
"unsloth_max_tool_calls_per_message",
"unsloth_tool_call_timeout",
"unsloth_chat_inference_params",
@ -409,7 +411,10 @@ export function GeneralTab() {
description: t("settings.general.rag.reindexWarning"),
});
} catch (error) {
if (error instanceof EmbeddingModelVerificationError) {
// A hard security block cannot be forced; keep the "save anyway" action hidden.
if (error instanceof EmbeddingModelBlockedError) {
setEmbeddingModelNeedsForce(false);
} else if (error instanceof EmbeddingModelVerificationError) {
setEmbeddingModelNeedsForce(true);
}
setEmbeddingModelError(

23
studio/frontend/src/features/settings/tabs/resources-tab.tsx
View file

@ -47,6 +47,15 @@ function formatGb(value: number | null | undefined): string {
return `${safe.toFixed(digits)} GB`;
}
// RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching
// nvidia-smi and PyTorch, so label those readouts GiB. Disk stays on formatGb
// because the backend reports disk in decimal GB (bytes / 1e9).
function formatGiB(value: number | null | undefined): string {
const safe = isFiniteNumber(value) ? Math.max(0, value) : 0;
const digits = safe >= 10 ? 1 : 2;
return `${safe.toFixed(digits)} GiB`;
}
function formatMb(value: number | null | undefined): string {
const safe = isFiniteNumber(value) ? Math.max(0, value) : 0;
return `${Math.round(safe).toLocaleString()} MB`;
@ -300,9 +309,9 @@ export function ResourcesTab() {
/>
<MetricTile
label={t("settings.resources.liveMonitor.ram")}
value={`${formatGb(metrics.ramUsed)} / ${formatGb(metrics.ramTotal)}`}
value={`${formatGiB(metrics.ramUsed)} / ${formatGiB(metrics.ramTotal)}`}
detail={t("settings.resources.liveMonitor.free", {
value: formatGb(systemInfo.memory?.available_gb),
value: formatGiB(systemInfo.memory?.available_gb),
})}
percent={systemInfo.memory?.percent_used ?? 0}
/>
@ -318,13 +327,13 @@ export function ResourcesTab() {
label={t("settings.resources.liveMonitor.vram")}
value={
hasGpu
? `${formatGb(metrics.vramUsed)} / ${formatGb(metrics.vramTotal)}`
? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}`
: t("settings.resources.liveMonitor.noGpu")
}
detail={
hasGpu
? t("settings.resources.liveMonitor.free", {
value: formatGb(metrics.vramFree),
value: formatGiB(metrics.vramFree),
})
: backendLabel
}
@ -373,17 +382,17 @@ export function ResourcesTab() {
<div className="grid gap-1 text-xs text-muted-foreground sm:grid-cols-3 sm:gap-2">
<span className="min-w-0 truncate font-mono tabular-nums">
{t("settings.resources.gpu.used", {
value: formatGb(used),
value: formatGiB(used),
})}
</span>
<span className="min-w-0 truncate font-mono tabular-nums sm:text-center">
{t("settings.resources.gpu.free", {
value: formatGb(free),
value: formatGiB(free),
})}
</span>
<span className="min-w-0 truncate font-mono tabular-nums sm:text-right">
{t("settings.resources.gpu.total", {
value: formatGb(total),
value: formatGiB(total),
})}
</span>
</div>

4
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -411,7 +411,7 @@ function LiveGpuPanel({
value={index}
className="bg-popover text-popover-foreground dark:bg-zinc-900 dark:text-zinc-100"
>
GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GB` : "N/A"})
GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GiB` : "N/A"})
</option>
))}
</select>
@ -446,7 +446,7 @@ function LiveGpuPanel({
icon={<HugeiconsIcon icon={RamMemoryIcon} className="size-3.5" />}
value={
currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null
? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB`
? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GiB`
: "--"
}
pct={currentGpu.vram_utilization_pct ?? 0}

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

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

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

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

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

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

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

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

8
studio/install_python_stack.py
View file

@ -1538,6 +1538,10 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = (
)
LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
# mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the
# QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242.
MLX_LM_BAD_VERSION_EXCLUSION = "!=0.31.3"
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides).
# _uv_safe_path: uv truncates UV_OVERRIDE at the first space too (issue #6503).
_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
@ -2104,6 +2108,8 @@ def install_python_stack() -> int:
# macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the
# mlx-vlm / mlx-lm transformers pin -- set at module load).
# Exclude mlx-lm 0.31.3 (see MLX_LM_BAD_VERSION_EXCLUSION); it broke
# gemma4 / qwen3_5 QK-norm loading. mlx-lm #1242.
if IS_MAC_ARM and not skip_base:
_progress("MLX stack (Apple Silicon)")
pip_install(
@ -2112,7 +2118,7 @@ def install_python_stack() -> int:
"--upgrade",
"mlx",
"mlx-metal",
"mlx-lm",
f"mlx-lm{MLX_LM_BAD_VERSION_EXCLUSION}",
"mlx-vlm",
)

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