diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index a1e7b2efa6..33ac3b9bd8 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -72,6 +72,31 @@ concurrency: permissions: contents: read +# ────────────────────────────────────────────────────────────────────── +# Network-resilience knobs, applied to every job/step. These add retries +# and backoff ONLY; they do not relax a single integrity check. cargo +# still resolves against Cargo.lock (--locked), pip still verifies the +# wheels it downloads, npm still enforces package-lock integrity, the +# harden-runner egress allowlists below are unchanged, and every action +# stays SHA-pinned. The advisory-audit run on 2026-05-29 red-failed when +# one crates.io tarball fetch hit "Recv failure: Connection reset by +# peer" (curl 56); cargo's default of 3 retries over an HTTP/2-multiplexed +# connection did not recover. The settings below make that class of +# transient fault self-heal instead of failing the whole run. +env: + # pip: raise the built-in retry count and per-connection timeout. + PIP_RETRIES: "10" + PIP_DEFAULT_TIMEOUT: "60" + # cargo: retry network ops and disable HTTP/2 multiplexing -- the + # documented mitigation for the curl-56 connection resets above. + CARGO_NET_RETRY: "10" + CARGO_HTTP_MULTIPLEXING: "false" + CARGO_NET_GIT_FETCH_WITH_CLI: "true" + # npm: retry registry fetches with capped exponential backoff. + NPM_CONFIG_FETCH_RETRIES: "5" + NPM_CONFIG_FETCH_RETRY_MINTIMEOUT: "2000" + NPM_CONFIG_FETCH_RETRY_MAXTIMEOUT: "60000" + jobs: # ───────────────────────────────────────────────────────────────────── # Combined advisory-DB audit: pip-audit + npm audit + cargo audit @@ -140,7 +165,7 @@ jobs: - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 - - uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2.9.1 + - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 with: workspaces: studio/src-tauri -> target @@ -153,8 +178,23 @@ jobs: # crashes with a TOML parse error on that file. # npm audit is bundled with the node toolchain, no install. run: | - python -m pip install --upgrade pip 'pip-audit>=2.7' - cargo install --locked --version '^0.22' cargo-audit + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + retry 5 python -m pip install --upgrade pip 'pip-audit>=2.7' + # --locked keeps the resolved tree identical to Cargo.lock; the + # CARGO_NET_* env above plus this outer loop survive transient + # crates.io connection resets without weakening that guarantee. + retry 5 cargo install --locked --version '^0.22' cargo-audit # ───────────────────────────────────────────────────────────── # Python: pip-audit @@ -330,32 +370,60 @@ jobs: # ───────────────────────────────────────────────────────────── # OSV-Scanner: cross-ecosystem advisory DB (PyPI + npm + cargo) # ───────────────────────────────────────────────────────────── + - name: Download + verify OSV-Scanner + # Split out from the scan below so binary integrity is a HARD gate: + # a checksum mismatch (swapped release asset, the Trivy-style pivot + # this workflow refuses) fails the job instead of being swallowed by + # the scan step's continue-on-error. A download still failing after + # retries is transient, so we skip the scan rather than red-fail. + # SHA-256 verified BEFORE chmod +x / exec. Bump OSV_SHA256 in lockstep + # with OSV_VERSION (value from the release's osv-scanner_SHA256SUMS). + run: | + set -euo pipefail + OSV_VERSION="v2.0.2" + OSV_SHA256="3abcfd7126c453a00421487e721b296e0cb68085bd431d6cef60872774170fc8" + if ! curl --proto '=https' --tlsv1.2 -fsSL \ + --retry 5 --retry-delay 3 --retry-connrefused --retry-all-errors \ + -o /tmp/osv-scanner \ + "https://github.com/google/osv-scanner/releases/download/${OSV_VERSION}/osv-scanner_linux_amd64"; then + echo "::warning::osv-scanner download failed after retries; skipping scan" >&2 + rm -f /tmp/osv-scanner + exit 0 # transient availability: do not red-fail the job + fi + if ! echo "${OSV_SHA256} /tmp/osv-scanner" | sha256sum -c -; then + echo "::error::osv-scanner checksum mismatch; refusing to execute" >&2 + rm -f /tmp/osv-scanner + exit 1 # integrity failure: hard-fail + fi + chmod +x /tmp/osv-scanner + /tmp/osv-scanner --version + - name: OSV-Scanner (PyPI + npm + cargo, cross-ecosystem advisories) # OSV's advisory feed is a superset of GitHub-Advisory + RustSec # + npm advisories; running it alongside the per-ecosystem audit # tools catches CVEs that haven't propagated to the per-ecosystem # DBs yet (e.g. langchain-core CVE-2025-68664 was on OSV before # GitHub Advisory). Single binary, one transitive resolver, all - # three lockfile types in one pass. Non-blocking until baselines - # close. + # three lockfile types in one pass. Binary is checksum-verified in + # the step above; only the advisory scan stays non-blocking until + # baselines close. continue-on-error: true run: | set +e - # OSV-Scanner ships a raw binary (no tarball) in v2.x. - curl -fsSL -o /tmp/osv-scanner \ - https://github.com/google/osv-scanner/releases/download/v2.0.2/osv-scanner_linux_amd64 - chmod +x /tmp/osv-scanner - /tmp/osv-scanner --version - /tmp/osv-scanner scan source \ - --lockfile=studio/frontend/package-lock.json \ - --lockfile=studio/src-tauri/Cargo.lock \ - --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ - --lockfile=requirements.txt:audit-reqs/studio.txt \ - --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ - --lockfile=requirements.txt:audit-reqs/overrides.txt \ - --lockfile=requirements.txt:audit-reqs/extras.txt \ - --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ - --format=table 2>&1 | tee logs-osv-scanner.txt + if [ ! -x /tmp/osv-scanner ]; then + echo "osv-scanner unavailable this run; skipping scan" | tee logs-osv-scanner.txt + else + /tmp/osv-scanner scan source \ + --lockfile=studio/frontend/package-lock.json \ + --lockfile=studio/src-tauri/Cargo.lock \ + --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ + --lockfile=requirements.txt:audit-reqs/studio.txt \ + --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ + --lockfile=requirements.txt:audit-reqs/overrides.txt \ + --lockfile=requirements.txt:audit-reqs/extras.txt \ + --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ + --format=table 2>&1 | tee logs-osv-scanner.txt + fi { echo "## OSV-Scanner (cross-ecosystem)" echo @@ -1075,7 +1143,23 @@ jobs: # new-install-script gate below protects against, and we must # not run any third-party hook to set up the audit. working-directory: studio/frontend - run: npm ci --ignore-scripts + run: | + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + # --ignore-scripts is mandatory here (no third-party hook runs); + # the retry only re-attempts the registry fetch, it never relaxes + # that flag or the package-lock integrity check npm ci enforces. + retry 5 npm ci --ignore-scripts - name: npm audit signatures (informational) # Surfaces unsigned / mis-signed packages from the npm diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index 6def56f769..c1297b84cc 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -296,6 +296,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: openai-anthropic-log @@ -771,6 +773,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: tool-calling-log @@ -1043,6 +1047,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: json-images-log diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 1096b1abb4..e08ac6ca68 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -289,6 +289,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: openai-anthropic-log @@ -649,6 +651,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: tool-calling-log @@ -1025,6 +1029,8 @@ jobs: - name: Upload logs # Always upload so green runs are still reviewable. if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: json-images-log diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 6253b9d213..bbe6b9e33a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -65,6 +65,18 @@ jobs: with: persist-credentials: false + # Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit + # test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke. + - name: setup.ps1 unit test (Resolve-CudaToolkit) + shell: pwsh + run: | + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "setup.ps1 parsed with no errors" + pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' @@ -360,6 +372,9 @@ jobs: - name: Collect llama-server logs if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true shell: bash # Copy llama-server's own stdout/stderr (teed by Studio under # ~/.unsloth/studio/logs/llama-server/) into the workspace so @@ -373,6 +388,8 @@ jobs: - name: Upload logs if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: windows-openai-anthropic-log @@ -788,6 +805,9 @@ jobs: - name: Collect llama-server logs if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true shell: bash # Copy llama-server's own stdout/stderr (teed by Studio under # ~/.unsloth/studio/logs/llama-server/) into the workspace so @@ -801,6 +821,8 @@ jobs: - name: Upload logs if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: windows-tool-calling-log @@ -1186,6 +1208,9 @@ jobs: - name: Collect llama-server logs if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true shell: bash # Copy llama-server's own stdout/stderr (teed by Studio under # ~/.unsloth/studio/logs/llama-server/) into the workspace so @@ -1199,6 +1224,8 @@ jobs: - name: Upload logs if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: windows-json-images-log diff --git a/README.md b/README.md index 948d84a789..562c35ff1a 100644 --- a/README.md +++ b/README.md @@ -72,10 +72,13 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**. ```bash curl -fsSL https://unsloth.ai/install.sh | sh ``` +Use the same command to update. + #### Windows: ```powershell irm https://unsloth.ai/install.ps1 | iex ``` +Use the same command to update. #### Launch ```bash @@ -83,9 +86,6 @@ unsloth studio -p 8888 ``` For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally. -#### Update -To update, use the same install commands above or use `unsloth studio update`. - #### Docker Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run: ```bash @@ -171,7 +171,9 @@ unsloth studio -p 8888 ``` Then to update : ```bash -unsloth studio update +cd unsloth && git pull +./install.sh --local +unsloth studio -p 8888 ``` #### Developer installs: Windows PowerShell: @@ -184,7 +186,9 @@ unsloth studio -p 8888 ``` Then to update : ```bash -unsloth studio update +cd unsloth && git pull +./install.sh --local +unsloth studio -p 8888 ``` #### Nightly: MacOS, Linux, WSL: diff --git a/install.ps1 b/install.ps1 index cab66f5ae1..47c72bcdc1 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1566,7 +1566,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.10" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1580,7 +1580,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.10" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1627,7 +1627,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.10" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1639,7 +1639,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.10" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1667,7 +1667,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.10" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index f0af60e2d6..15755c589f 100755 --- a/install.sh +++ b/install.sh @@ -1530,17 +1530,51 @@ if [ -x "$VENV_DIR/bin/python" ]; then : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true fi -# Guard against Python 3.13.8 torch import bug on Apple Silicon -# (skip when the user explicitly chose a version via --python) +# Guard against two independent Apple Silicon venv problems, in order: +# 1. uv may create the venv from a cached x86_64 (Rosetta) Python when a +# same-version x86_64 build is already cached (often because uv itself +# is an x86_64 build). That venv reports x86_64 to wheel resolvers, and +# PyTorch ships no macOS wheels on the CPU index for any architecture, +# so the torch install can never resolve. Recreate it with an +# arch-explicit arm64 CPython. +# 2. Python 3.13.8 has a known torch import bug. +# The two are independent: a venv may be x86_64 and, once recreated, still +# land on 3.13.8. So we re-inspect the interpreter between the checks instead +# of chaining them with elif, guaranteeing both invariants hold on whatever +# venv we end up with. Skip both when the user explicitly chose an interpreter +# via --python. if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then - _PY_VER=$("$VENV_DIR/bin/python" -c \ - "import sys; print('{}.{}.{}'.format(*sys.version_info[:3]))" 2>/dev/null || echo "") + _inspect_venv() { + "$VENV_DIR/bin/python" -c \ + "import platform, sys; print(platform.machine(), '{}.{}.{}'.format(*sys.version_info[:3]))" \ + 2>/dev/null || echo " " + } + _info=$(_inspect_venv) + _VENV_ARCH=${_info%% *} + _PY_VER=${_info##* } + + if [ "$_VENV_ARCH" = "x86_64" ]; then + echo " WARNING: venv was created with an x86_64 (Rosetta) Python on Apple Silicon." + echo " Recreating venv with native arm64 Python ${PYTHON_VERSION}..." + rm -rf "$VENV_DIR" + run_install_cmd "recreate venv (arm64)" uv venv "$VENV_DIR" \ + --python "cpython-${PYTHON_VERSION}-macos-aarch64-none" + if [ -x "$VENV_DIR/bin/python" ]; then + : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true + fi + # Re-inspect: the recreated arm64 venv may still be 3.13.8. + _info=$(_inspect_venv) + _VENV_ARCH=${_info%% *} + _PY_VER=${_info##* } + fi + if [ "$_PY_VER" = "3.13.8" ]; then echo " WARNING: Python 3.13.8 has a known torch import bug." echo " Recreating venv with Python 3.12..." rm -rf "$VENV_DIR" PYTHON_VERSION="3.12" - run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + run_install_cmd "recreate venv" uv venv "$VENV_DIR" \ + --python "cpython-${PYTHON_VERSION}-macos-aarch64-none" if [ -x "$VENV_DIR/bin/python" ]; then : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true fi @@ -2049,7 +2083,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.8" unsloth-zoo + "unsloth>=2026.5.10" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2062,7 +2096,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.8" unsloth-zoo + "unsloth>=2026.5.10" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2266,7 +2300,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.5.8" unsloth-zoo + "unsloth>=2026.5.10" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2284,7 +2318,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.10" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2316,7 +2350,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.10" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." diff --git a/pyproject.toml b/pyproject.toml index aef88d90f5..acc65f12cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "wheel>=0.42.0", "packaging", "numpy", @@ -90,7 +90,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "torchvision", "unsloth[triton]", ] @@ -580,7 +580,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index b4ec0ccd94..85b567885f 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -176,12 +176,21 @@ def build_mcp_providers( ) -> list: from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] + # Same gate as the chat MCP path: stdio providers spawn a local subprocess, + # so only build them when this host allows it (desktop / explicit opt-in). + # Skip them otherwise so a recipe carried onto a hosted host cannot spawn. + from core.inference.mcp_client import stdio_mcp_enabled + + stdio_allowed = stdio_mcp_enabled() + providers: list[MCPProvider | LocalStdioMCPProvider] = [] for provider in recipe.get("mcp_providers", []): if not isinstance(provider, dict): continue provider_type = provider.get("provider_type") if provider_type == "stdio": + if not stdio_allowed: + continue env = provider.get("env") if not isinstance(env, dict): env = {} diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index bc792c3b99..cdb0fdebff 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -218,6 +218,8 @@ class AnthropicStreamEmitter: def __init__(self) -> None: self.block_index: int = 0 self._text_block_open: bool = False + self._open_tool_call_id: Optional[str] = None + self._open_tool_args_sent: bool = False self._prev_text: str = "" self._usage: dict = {} @@ -263,8 +265,10 @@ class AnthropicStreamEmitter: def finish(self, stop_reason: str = "end_turn") -> list[str]: """Close any open block and emit message_delta + message_stop.""" events = [] - if self._text_block_open: + if self._text_block_open or self._open_tool_call_id is not None: events.append(self._close_block()) + self._open_tool_call_id = None + self._open_tool_args_sent = False events.append( build_anthropic_sse_event( "message_delta", @@ -310,12 +314,26 @@ class AnthropicStreamEmitter: return events def _handle_tool_start(self, event: dict) -> list[str]: + tool_call_id = event.get("tool_call_id", "") + args = event.get("arguments", {}) + if tool_call_id and self._open_tool_call_id == tool_call_id: + return self._tool_arguments_delta(args) + events = [] - # Close current text block if open + # Close current text block if open. if self._text_block_open: events.append(self._close_block()) - # Open a tool_use block + # Defensive: if a replacement/different tool_start arrives while a + # tool_use block is open, close the stale block before starting another. + elif self._open_tool_call_id is not None: + events.append(self._close_block()) + self._open_tool_call_id = None + self._open_tool_args_sent = False + + # Open a tool_use block. self.block_index += 1 + self._open_tool_call_id = tool_call_id + self._open_tool_args_sent = False events.append( build_anthropic_sse_event( "content_block_start", @@ -324,35 +342,43 @@ class AnthropicStreamEmitter: "index": self.block_index, "content_block": { "type": "tool_use", - "id": event.get("tool_call_id", ""), + "id": tool_call_id, "name": event.get("tool_name", ""), "input": {}, }, }, ) ) - # Emit the arguments as input_json_delta - args = event.get("arguments", {}) - if args: - events.append( - build_anthropic_sse_event( - "content_block_delta", - { - "type": "content_block_delta", - "index": self.block_index, - "delta": { - "type": "input_json_delta", - "partial_json": json.dumps(args), - }, - }, - ) - ) + events.extend(self._tool_arguments_delta(args)) return events + def _tool_arguments_delta(self, args: dict) -> list[str]: + if not args: + return [] + if self._open_tool_args_sent: + return [] + self._open_tool_args_sent = True + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": json.dumps(args), + }, + }, + ) + ] + def _handle_tool_end(self, event: dict) -> list[str]: events = [] - # Close the tool_use block - events.append(self._close_block()) + # Close the tool_use block. + if self._open_tool_call_id is not None or self._text_block_open: + events.append(self._close_block()) + self._open_tool_call_id = None + self._open_tool_args_sent = False # Emit custom tool_result event (non-standard, ignored by SDKs) events.append( build_anthropic_sse_event( diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b1aceb42cd..b5e0c029b6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -52,6 +52,7 @@ from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) from core.inference.tool_call_parser import ( + RENDER_HTML_REPEAT_NUDGE, parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, ) @@ -2604,6 +2605,105 @@ class LlamaCppBackend: # ── Lifecycle ───────────────────────────────────────────────── + # GGUF ``general.architecture`` values for diffusion / image models. + # llama.cpp proper has no such architectures, so loading one as a chat + # model dies with "unknown model architecture: ''". These match + # the patched stable-diffusion.cpp / ComfyUI-GGUF enums (LLM_ARCH_FLUX, + # LLM_ARCH_QWEN_IMAGE, ...). Unsloth publishes FLUX and Qwen-Image GGUFs + # under https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs. + # Matched exactly (not as a substring) so a chat arch merely containing a + # short token like "wan"/"sd1" (e.g. "taiwan") is not misrouted to Images. + _DIFFUSION_ARCHES = frozenset( + ( + "qwen_image", + "flux", + "sd1", + "sdxl", + "sd3", + "aura", + "hidream", + "cosmos", + "ltxv", + "hyvid", + "wan", + "lumina2", + ) + ) + + @staticmethod + def _classify_llama_start_failure( + output: str, + gguf_path: Optional[str], + model_identifier: Optional[str], + ) -> str: + """Explain *why* llama-server failed to start, from its output. + + Several distinct failures all otherwise collapse into the same + opaque "invalid GGUF or out of memory" message. The worst case is + a diffusion / image GGUF (FLUX, Qwen-Image, ...) loaded as a chat + model: the file is perfectly valid and there is plenty of memory, + but llama.cpp has no such architecture, so the user is told to free + memory that was never the problem (issue #5842). Pick the most + specific message the captured output supports. + """ + lowered = (output or "").lower() + + # Detect Ollama source up front so the arch branch can keep the + # Ollama hint instead of the generic "unsupported arch" message. + gguf = gguf_path or "" + is_ollama = ( + ".studio_links" in gguf + or os.sep + "ollama_links" + os.sep in gguf + or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf + or (model_identifier or "").startswith("ollama/") + ) + + # "unknown model architecture: ''": diffusion -> Images page, + # Ollama -> Ollama hint, else a precise "unsupported" message. Exact + # match so chat archs are never misrouted. + arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered) + if arch_match: + arch = arch_match.group(1) + if arch in LlamaCppBackend._DIFFUSION_ARCHES: + return ( + f"'{arch}' is a diffusion (image-generation) GGUF, which " + "llama-server cannot run as a chat/completion model. Use " + "Studio's Images page to generate with local diffusion " + "GGUFs such as FLUX and Qwen-Image." + ) + if is_ollama: + return ( + "Some Ollama models do not work with llama.cpp. Try a " + "different model, or use this model directly through " + "Ollama instead." + ) + return ( + f"llama.cpp does not support this GGUF's model architecture " + f"('{arch}'). The file is valid, but this model type cannot " + "be run with llama-server." + ) + + # Other Ollama compat failures that do not name an arch. Only when + # the output shows a GGUF compat issue, not OOM / missing binaries. + if is_ollama: + gguf_compat_hints = ( + "key not found", + "unknown model architecture", + "failed to load model", + ) + if any(h in lowered for h in gguf_compat_hints): + return ( + "Some Ollama models do not work with llama.cpp. Try a " + "different model, or use this model directly through " + "Ollama instead." + ) + + # Fallback: genuinely unknown failure (OOM, missing binary, ...). + return ( + "llama-server failed to start. " + "Check that the GGUF file is valid and you have enough memory." + ) + def load_model( self, *, @@ -3396,31 +3496,12 @@ class LlamaCppBackend: # Wait for llama-server to become healthy if not self._wait_for_health(timeout = 600.0): self._kill_process() - _gguf = gguf_path or "" - _is_ollama = ( - ".studio_links" in _gguf - or os.sep + "ollama_links" + os.sep in _gguf - or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf - or (self._model_identifier or "").startswith("ollama/") - ) - # Only show the Ollama-specific message when the server - # output indicates a GGUF compatibility issue, not for - # unrelated failures like OOM or missing binaries. - if _is_ollama: - _output = "\n".join(self._stdout_lines[-50:]).lower() - _gguf_compat_hints = ( - "key not found", - "unknown model architecture", - "failed to load model", - ) - if any(h in _output for h in _gguf_compat_hints): - raise RuntimeError( - "Some Ollama models do not work with llama.cpp. " - "Try a different model, or use this model directly through Ollama instead." - ) raise RuntimeError( - "llama-server failed to start. " - "Check that the GGUF file is valid and you have enough memory." + self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + self._model_identifier, + ) ) self._healthy = True @@ -4553,6 +4634,7 @@ class LlamaCppBackend: # a transient failure are allowed (only block when the previous # identical call succeeded). _tool_call_history: list[tuple[str, bool]] = [] # (key, failed) + _render_html_succeeded = False # ── Re-prompt on plan-without-action ───────────────── # When the model describes what it intends to do (forward-looking @@ -4627,6 +4709,7 @@ class LlamaCppBackend: _iter_timings = None _stream_done = False _last_emitted = "" + provisional_render_html_tool_call_ids = set() stream_timeout = httpx.Timeout( connect = 10, @@ -4736,6 +4819,33 @@ class LlamaCppBackend: tool_calls_acc[idx]["function"][ "arguments" ] += func["arguments"] + current_name = tool_calls_acc[idx][ + "function" + ].get("name", "") + fallback_id = f"call_{idx}" + current_id = tool_calls_acc[idx].get( + "id", fallback_id + ) + already_started = ( + current_id + in provisional_render_html_tool_call_ids + ) + has_real_id = current_id != fallback_id + if ( + current_name == "render_html" + and not _render_html_succeeded + and not already_started + and has_real_id + ): + provisional_render_html_tool_call_ids.add( + current_id + ) + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": current_id, + "arguments": {}, + } continue # ── Reasoning tokens ── @@ -4917,13 +5027,25 @@ class LlamaCppBackend: "content": _stripped, } ) + available_tool_names = [ + tool.get("function", {}).get("name") + for tool in tools + if isinstance(tool, dict) + and isinstance(tool.get("function"), dict) + ] + available_tool_names = [ + name for name in available_tool_names if name + ] + tool_hint = ( + " or ".join(available_tool_names) or "an available tool" + ) conversation.append( { "role": "user", "content": ( "STOP. Do NOT write code or explain. " "You MUST call a tool NOW. " - "Call web_search or python immediately." + f"Call {tool_hint} immediately." ), } ) @@ -5095,7 +5217,12 @@ class LlamaCppBackend: arguments = json.loads(raw_args) except (json.JSONDecodeError, ValueError): if auto_heal_tool_calls: - arguments = {"query": raw_args} + heal_key = { + "python": "code", + "terminal": "command", + "render_html": "code", + }.get(tool_name, "query") + arguments = {heal_key: raw_args} else: arguments = {"raw": raw_args} else: @@ -5132,14 +5259,18 @@ class LlamaCppBackend: ) else: status_text = f"Calling: {tool_name}" - yield {"type": "status", "text": status_text} + _repeat_render_html = ( + tool_name == "render_html" and _render_html_succeeded + ) + if not _repeat_render_html: + yield {"type": "status", "text": status_text} - yield { - "type": "tool_start", - "tool_name": tool_name, - "tool_call_id": tc.get("id", ""), - "arguments": arguments, - } + yield { + "type": "tool_start", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "arguments": arguments, + } # ── Duplicate call detection ────────────── # str(dict) is stable here: arguments always comes from @@ -5147,7 +5278,9 @@ class LlamaCppBackend: # so insertion order is deterministic (Python 3.7+). _tc_key = tool_name + str(arguments) _prev = _tool_call_history[-1] if _tool_call_history else None - if _prev and _prev[0] == _tc_key and not _prev[1]: + if _repeat_render_html: + result = RENDER_HTML_REPEAT_NUDGE + elif _prev and _prev[0] == _tc_key and not _prev[1]: result = ( "You already made this exact call. " "Do not repeat the same tool call. " @@ -5186,12 +5319,13 @@ class LlamaCppBackend: tool_context = tool_context, ) - yield { - "type": "tool_end", - "tool_name": tool_name, - "tool_call_id": tc.get("id", ""), - "result": result, - } + if not _repeat_render_html: + yield { + "type": "tool_end", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "result": result, + } # Nudge model to try a different approach on errors _error_prefixes = ( @@ -5207,6 +5341,8 @@ class LlamaCppBackend: _is_error = isinstance(result, str) and result.lstrip().startswith( _error_prefixes ) + if tool_name == "render_html" and not _is_error: + _render_html_succeeded = True _tool_call_history.append((_tc_key, _is_error)) # Strip image sentinel before feeding result to the LLM # (the full result with sentinel is still yielded via diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index a5e614899d..2ed1a630dc 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -5,6 +5,9 @@ from __future__ import annotations import asyncio import json +import os +import shlex +import sys from typing import Any, Optional from loggers import get_logger @@ -16,7 +19,55 @@ MCP_TOOL_PREFIX = "mcp__" _oauth_token_store = None +def is_stdio(address: str) -> bool: + """A non-HTTP address is a local stdio command, e.g. + 'npx -y @modelcontextprotocol/server-filesystem /path'.""" + return not address.strip().lower().startswith(("http://", "https://")) + + +def parse_stdio_command(address: str) -> list[str]: + """Split a stdio command line into argv. Shared by route validation and the + transport so both agree on quoting (notably Windows backslash paths).""" + posix = sys.platform != "win32" + parts = shlex.split(address, posix = posix) + if not posix: + # posix=False keeps backslash paths intact but also keeps the surrounding + # quotes on a token. Strip a matched pair so the argv reaches the + # subprocess clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node). + parts = [ + p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p + for p in parts + ] + return parts + + +def stdio_mcp_enabled() -> bool: + """stdio MCP servers spawn local processes as the backend user (and bypass + the python/terminal sandbox), so they are only allowed when the backend + host is the user's own machine. The Tauri desktop app sets + UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see main.py); advanced localhost / + self-hosted users can opt in with the same variable. It stays off for + Colab and any network (0.0.0.0) bind.""" + return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1" + + +# Probe timeouts for discovering a server's tool list. OAuth needs minutes for +# first-connect/expired-token browser sign-in; stdio allows for first-run +# package download (e.g. `npx -y ...`); HTTP fails fast. +_HTTP_PROBE_TIMEOUT = 8.0 +_OAUTH_PROBE_TIMEOUT = 305.0 +_STDIO_PROBE_TIMEOUT = 60.0 + + +def probe_timeout(address: str, use_oauth: bool) -> float: + if use_oauth: + return _OAUTH_PROBE_TIMEOUT + return _STDIO_PROBE_TIMEOUT if is_stdio(address) else _HTTP_PROBE_TIMEOUT + + def parse_server_headers(server: dict) -> Optional[dict]: + """Parsed headers_json. For stdio servers this dict is the process + environment instead of HTTP headers (see _client).""" raw = server.get("headers_json") if not raw: return None @@ -63,6 +114,28 @@ async def clear_oauth_tokens_async(url: str) -> None: def _client(url: str, headers: Optional[dict], use_oauth: bool = False): from fastmcp import Client + + if is_stdio(url): + # Belt-and-suspenders: never spawn unless stdio is enabled on this host. + if not stdio_mcp_enabled(): + raise PermissionError("stdio MCP servers are disabled on this host") + from fastmcp.client.transports import StdioTransport + + parts = parse_stdio_command(url) + if not parts: + raise ValueError(f"Empty stdio command: {url!r}") + # env vars ride the headers field (merged over the SDK's safe default env). + # keep_alive=False tears the subprocess down on exit, so a one-shot + # probe/tool call never leaves an orphan process. + return Client( + StdioTransport( + command = parts[0], + args = parts[1:], + env = headers or None, + keep_alive = False, + ) + ) + from fastmcp.client.transports import SSETransport, StreamableHttpTransport from fastmcp.mcp_config import infer_transport_type_from_url diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index edfae5ce16..9487eead3d 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -18,6 +18,7 @@ cumulative text and dispatches them via ``core.inference.tools``. """ import json +import re import threading from typing import Callable, Generator, Optional from urllib.parse import urlparse @@ -27,6 +28,7 @@ from loggers import get_logger from core.inference.tool_call_parser import ( BUDGET_EXHAUSTED_NUDGE, DUPLICATE_CALL_NUDGE, + RENDER_HTML_REPEAT_NUDGE, TOOL_ERROR_NUDGE, TOOL_ERROR_PREFIXES, TOOL_XML_SIGNALS, @@ -66,7 +68,34 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return f"Calling: {tool_name}" -_CANONICAL_HEAL_ARG = {"python": "code", "terminal": "command"} +_CANONICAL_HEAL_ARG = { + "python": "code", + "terminal": "command", + "render_html": "code", +} + + +_FUNCTION_SIGNAL_RE = re.compile(r"") +_TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') + + +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("") + if not function_match and tool_call_index < 0: + 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 def _coerce_arguments(raw_args, *, heal: bool, tool_name: str = "") -> dict: @@ -136,6 +165,7 @@ def run_safetensors_tool_loop( """ conversation = list(messages) tool_call_history: list[tuple[str, bool]] = [] + render_html_succeeded = False final_attempt_done = False allowed_tool_names = { (tool.get("function") or {}).get("name") @@ -162,6 +192,8 @@ def run_safetensors_tool_loop( content_accum = "" cumulative_display = "" last_emitted = "" + provisional_render_html_started = False + provisional_render_html_id = f"call_{next_call_id}" gen = single_turn(conversation) prev_cumulative = "" @@ -180,6 +212,18 @@ def run_safetensors_tool_loop( content_accum += delta if detect_state == _state_draining: + if ( + not render_html_succeeded + and not provisional_render_html_started + and _detect_render_html_tool_start(content_accum) + ): + provisional_render_html_started = True + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": provisional_render_html_id, + "arguments": {}, + } continue if detect_state == _state_streaming: @@ -197,6 +241,18 @@ def run_safetensors_tool_loop( yield {"type": "content", "text": cleaned_before} cumulative_display = candidate detect_state = _state_draining + if ( + not render_html_succeeded + and not provisional_render_html_started + and _detect_render_html_tool_start(content_accum) + ): + provisional_render_html_started = True + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": provisional_render_html_id, + "arguments": {}, + } continue cumulative_display = candidate cleaned = strip_tool_markup(cumulative_display) @@ -223,6 +279,18 @@ def run_safetensors_tool_loop( if is_match: detect_state = _state_draining + if ( + not render_html_succeeded + and not provisional_render_html_started + and _detect_render_html_tool_start(content_accum) + ): + provisional_render_html_started = True + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": provisional_render_html_id, + "arguments": {}, + } elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS: continue else: @@ -283,6 +351,13 @@ def run_safetensors_tool_loop( # literal "" prose is preserved. if content_accum: yield {"type": "content", "text": content_accum} + if provisional_render_html_started: + yield { + "type": "tool_end", + "tool_name": "render_html", + "tool_call_id": provisional_render_html_id, + "result": "Error: render_html tool call could not be parsed.", + } yield {"type": "status", "text": ""} return content_text = strip_tool_markup(content_accum, final = True) @@ -309,16 +384,20 @@ def run_safetensors_tool_loop( tool_name = tool_name, ) - yield {"type": "status", "text": _status_for_tool(tool_name, arguments)} - yield { - "type": "tool_start", - "tool_name": tool_name, - "tool_call_id": tc.get("id", ""), - "arguments": arguments, - } + repeat_render_html = tool_name == "render_html" and render_html_succeeded + if not repeat_render_html: + yield {"type": "status", "text": _status_for_tool(tool_name, arguments)} + yield { + "type": "tool_start", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "arguments": arguments, + } tc_key = tool_name + str(arguments) - if allowed_tool_names and tool_name not in allowed_tool_names: + if repeat_render_html: + result = RENDER_HTML_REPEAT_NUDGE + elif allowed_tool_names and tool_name not in allowed_tool_names: result = ( f"Error: tool '{tool_name}' is not enabled for this " "request. Use one of the enabled tools or provide a " @@ -347,16 +426,19 @@ def run_safetensors_tool_loop( logger.exception("Tool %s raised: %s", tool_name, exc) result = f"Error: tool raised an exception: {exc}" - yield { - "type": "tool_end", - "tool_name": tool_name, - "tool_call_id": tc.get("id", ""), - "result": result, - } + if not repeat_render_html: + yield { + "type": "tool_end", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "result": result, + } is_error = isinstance(result, str) and result.lstrip().startswith( TOOL_ERROR_PREFIXES ) + if tool_name == "render_html" and not is_error: + render_html_succeeded = True tool_call_history.append((tc_key, is_error)) # Strip frontend image sentinel from the model's view. diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 2f94990623..dacbc19ac0 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -49,6 +49,12 @@ DUPLICATE_CALL_NUDGE = ( "provide your final answer now." ) +RENDER_HTML_REPEAT_NUDGE = ( + "Error: render_html was already called for this response. Do not call " + "render_html again in this response unless the user asks for changes. " + "Provide the final answer now." +) + TOOL_ERROR_NUDGE = ( "\n\nThe tool call encountered an issue. Please try a different " "approach or rephrase your request." @@ -70,6 +76,20 @@ _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") # `issue-number`, `repo-name`); using `\w+` here dropped those keys. _TC_PARAM_START_RE = re.compile(r"\s*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") +_PARAM_CLOSE_TAG = "" +_FUNC_CLOSE_TAG = "" + + +def _inside_open_parameter(content: str, pos: int) -> bool: + """Return True when ``pos`` falls inside an unclosed parameter value.""" + last_param_start = -1 + for match in _TC_PARAM_START_RE.finditer(content, 0, pos): + 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) def strip_tool_markup(text: str, *, final: bool = False) -> str: @@ -151,7 +171,11 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict # optional; don't use as body boundary because code # values can contain that literal. if not tool_calls: - func_starts = list(_TC_FUNC_START_RE.finditer(content)) + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) body_start = fm.end() diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 48d70aa67f..328f2a73a6 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -28,8 +28,11 @@ import urllib.request from core.inference.mcp_client import ( MCP_TOOL_PREFIX, call_tool_sync, + is_stdio, list_tools_async, parse_server_headers, + probe_timeout, + stdio_mcp_enabled, ) from storage import mcp_servers_db @@ -418,6 +421,35 @@ _workdirs: dict[str, str] = {} # Non-matching session_ids collapse to ``_invalid`` to block cross-session escapes. _SESSION_ID_RE = re.compile(r"\A[A-Za-z0-9_\-]{1,64}\Z") +_PROJECT_SESSION_PREFIX = "project-" + + +def _get_project_workdir(session_id: str) -> str | None: + if not session_id.startswith(_PROJECT_SESSION_PREFIX): + return None + project_id = session_id[len(_PROJECT_SESSION_PREFIX) :] + if not project_id or not _SESSION_ID_RE.match(project_id): + return None + try: + from storage.studio_db import ensure_chat_project_workspace + + project = ensure_chat_project_workspace(project_id) + except Exception: + logger.warning( + "Failed to resolve project sandbox for %s", session_id, exc_info = True + ) + return None + if not project: + return None + root_path = project.get("rootPath") + sandbox_path = project.get("sandboxPath") + if not root_path or not sandbox_path: + return None + root_real = os.path.realpath(root_path) + sandbox_real = os.path.realpath(sandbox_path) + if sandbox_real != root_real and not sandbox_real.startswith(root_real + os.sep): + return None + return sandbox_real def _get_workdir(session_id: str | None = None) -> str: @@ -427,7 +459,14 @@ def _get_workdir(session_id: str | None = None) -> str: if key not in _workdirs or not os.path.isdir(_workdirs[key]): home = os.path.expanduser("~") sandbox_root = os.path.join(home, "studio_sandbox") - if session_id and _SESSION_ID_RE.match(session_id): + project_workdir = ( + _get_project_workdir(session_id) + if session_id and _SESSION_ID_RE.match(session_id) + else None + ) + if project_workdir: + workdir = project_workdir + elif session_id and _SESSION_ID_RE.match(session_id): workdir = os.path.join(sandbox_root, session_id) if not os.path.realpath(workdir).startswith( os.path.realpath(sandbox_root) + os.sep @@ -450,6 +489,10 @@ def _get_workdir(session_id: str | None = None) -> str: return _workdirs[key] +def get_sandbox_workdir(session_id: str | None = None) -> str: + return _get_workdir(session_id) + + WEB_SEARCH_TOOL = { "type": "function", "function": { @@ -511,7 +554,6 @@ TERMINAL_TOOL = { }, } - # Lazy import: don't pull rag stack on inference paths that never see RAG. def _get_rag_tool_spec(): from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL @@ -519,7 +561,41 @@ def _get_rag_tool_spec(): return SEARCH_KNOWLEDGE_BASE_TOOL -ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()] +RENDER_HTML_TOOL = { + "type": "function", + "function": { + "name": "render_html", + "description": ( + "Render a self-contained HTML/CSS/JavaScript artifact for the user. " + "Call this at most once per assistant response unless the user " + "explicitly asks for changes in that response. Future user requests " + "for new artifacts may call render_html once. Put the entire document " + "in code, including any CSS in