diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 02595510d4..51f9e724d9 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -8,6 +8,9 @@ updates: groups: actions: patterns: ["*"] + actions-security: + applies-to: security-updates + patterns: ["*"] - package-ecosystem: "bun" directory: "/studio/frontend" @@ -16,6 +19,9 @@ updates: groups: bun-frontend: patterns: ["*"] + bun-frontend-security: + applies-to: security-updates + patterns: ["*"] - package-ecosystem: "npm" directory: "/studio/backend/core/data_recipe/oxc-validator" @@ -24,11 +30,12 @@ updates: groups: npm-oxc-validator: patterns: ["*"] + npm-oxc-validator-security: + applies-to: security-updates + patterns: ["*"] - # pip + cargo so security advisories on Python deps + the Tauri shell - # auto-generate PRs alongside the github-actions / bun / npm updates. - # Grouped weekly so we don't get one PR per dep; security advisories - # bypass the group and open immediately. + # pip + cargo grouped weekly; the *-security siblings batch + # advisories that would otherwise each open their own PR. - package-ecosystem: "pip" directory: "/" schedule: @@ -37,6 +44,9 @@ updates: groups: python: patterns: ["*"] + python-security: + applies-to: security-updates + patterns: ["*"] - package-ecosystem: "cargo" directory: "/studio/src-tauri" @@ -45,4 +55,21 @@ updates: groups: cargo-tauri: patterns: ["*"] + cargo-tauri-security: + applies-to: security-updates + patterns: ["*"] + + # bun owns version updates for /studio/frontend (above); GitHub + # fires npm-package advisories under npm_and_yarn, so this entry + # catches and groups them. limit: 0 suppresses version-update + # PRs, security updates flow through regardless. + - package-ecosystem: "npm" + directory: "/studio/frontend" + schedule: + interval: "weekly" + open-pull-requests-limit: 0 + groups: + npm-frontend-security: + applies-to: security-updates + patterns: ["*"] ... diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 4ad3d9f16a..0f6f89d354 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -234,9 +234,23 @@ jobs: # tests/conftest.py spoof which handles that. run: | set -euxo pipefail - git clone --depth=1 --branch="$UNSLOTH_ZOO_REF" \ - https://github.com/unslothai/unsloth-zoo \ - "$RUNNER_TEMP/unsloth-zoo" + # github.com occasionally 500s on the git fetch; retry so a + # single upstream blip does not fail CI. + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 --branch="$UNSLOTH_ZOO_REF" \ + https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps pip show unsloth_zoo @@ -2040,9 +2054,23 @@ jobs: # main-branch fixes flow into the smoke without a release). run: | set -euxo pipefail - git clone --depth=1 --branch="$UNSLOTH_ZOO_REF" \ - https://github.com/unslothai/unsloth-zoo \ - "$RUNNER_TEMP/unsloth-zoo" + # github.com occasionally 500s on the git fetch; retry so a + # single upstream blip does not fail CI. + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 --branch="$UNSLOTH_ZOO_REF" \ + https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps pip show unsloth_zoo diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 61e0566903..52038dd332 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -153,7 +153,20 @@ jobs: 'httpx==0.28.1' pip install --index-url https://download.pytorch.org/whl/cpu \ 'torch==2.10.0' - pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo" + # github.com occasionally 500s on the git fetch; retry the + # zoo install so a single upstream blip does not fail CI. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::pip install unsloth_zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::unsloth_zoo install failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done pip install -e . --no-deps # Real Apple Silicon sanity: confirm _IS_MLX activates on real diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 0fc8073e75..7ab1021ec9 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -57,6 +57,7 @@ on: - 'studio/src-tauri/Cargo.lock' - 'pyproject.toml' - 'scripts/scan_packages.py' + - 'scripts/scan_npm_packages.py' - '.github/workflows/security-audit.yml' push: branches: [main, pip] @@ -244,6 +245,27 @@ jobs: echo '```' } >> "$GITHUB_STEP_SUMMARY" + # ───────────────────────────────────────────────────────────── + # Pre-install lockfile supply-chain audit (npm + cargo). + # Catches structural anomalies (non-registry resolved URLs, + # missing integrity hashes, known IOC strings) BEFORE `npm + # audit` or OSV-Scanner consult the advisory DB. The advisory + # path is reactive -- there is a window between a malicious + # publication and the GHSA landing. This step fires on the + # injection pattern itself so it catches the same class of + # attack the moment the lockfile shape becomes wrong. + # ───────────────────────────────────────────────────────────── + - name: Lockfile supply-chain audit (pre-install scan) + run: | + python3 scripts/lockfile_supply_chain_audit.py + { + echo "## Lockfile supply-chain audit" + echo + echo "Scanned: studio/frontend/package-lock.json + studio/src-tauri/Cargo.lock" + echo + echo "No structural anomalies or known IOC strings." + } >> "$GITHUB_STEP_SUMMARY" + # ───────────────────────────────────────────────────────────── # npm: Studio frontend # ───────────────────────────────────────────────────────────── @@ -794,3 +816,82 @@ jobs: logs-scan-packages-${{ matrix.shard.id }}.txt audit-reqs/ retention-days: 30 + + # ───────────────────────────────────────────────────────────────────── + # npm: pre-install tarball content scan. + # ───────────────────────────────────────────────────────────────────── + npm-scan-packages: + # Counterpart to pip-scan-packages for the npm side. Reads + # studio/frontend/package-lock.json, downloads each resolved + # tarball DIRECTLY from registry.npmjs.org (never via `npm + # install` -- no lifecycle scripts ever run), verifies the + # lockfile integrity hash, unpacks each tarball into a sandboxed + # temp dir behind size / count / path-escape / symlink guards, + # and pattern-scans the extracted file contents for the + # signatures common to npm supply-chain attacks: + # + # - lifecycle (preinstall / install / postinstall / prepare) + # scripts in any package.json that fetch + execute external + # code, + # - C2 / exfiltration hosts (getsession.org, AWS IMDS, + # Kubernetes ServiceAccount token paths, GitHub Actions OIDC, + # HashiCorp Vault endpoints), + # - credential-stealing references (.npmrc, .aws/credentials, + # GITHUB_TOKEN / NPM_TOKEN in JS sources), + # - known IOC filenames (router_init.js, tanstack_runner.js, + # router_runtime.js), + # - obfuscation shapes (Function/eval against base64 blobs). + # + # Threat model: every tarball is hostile. Safety guarantees are + # documented at scripts/scan_npm_packages.py top-of-file. The + # script is stdlib-only so adding it does not increase the + # transitive supply-chain surface. + name: npm scan-packages (Studio frontend tarballs) + runs-on: ubuntu-latest + timeout-minutes: 30 + needs: [] + steps: + - name: Harden runner (egress audit) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: audit + disable-sudo: true + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Sanity-check scan_npm_packages.py + run: | + test -f scripts/scan_npm_packages.py + python3 -c "import ast; ast.parse(open('scripts/scan_npm_packages.py').read())" + + - name: Scan npm tarballs (declared + transitive, no install) + # The script exits 1 on HIGH/CRITICAL findings; we capture the + # full log and surface it in the step summary either way. It + # never runs `npm install`, never executes anything from a + # downloaded tarball, and only fetches from registry.npmjs.org. + # Initially non-blocking so the baseline can settle; drop + # continue-on-error once the baseline is clean for a week. + continue-on-error: true + run: | + set -o pipefail + LOG=logs-scan-npm.txt + python3 scripts/scan_npm_packages.py 2>&1 | tee "$LOG" + { + echo "## scan_npm_packages" + echo + echo '### Findings (tail)' + echo '```' + tail -300 "$LOG" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: always() + with: + name: scan-npm-packages-log + path: logs-scan-npm.txt + retention-days: 30 diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 4e8cc5c9c3..742cba9ed9 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -80,6 +80,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml index eb00e297a7..bde62c87f6 100644 --- a/.github/workflows/studio-frontend-ci.yml +++ b/.github/workflows/studio-frontend-ci.yml @@ -58,6 +58,14 @@ jobs: cache: 'npm' cache-dependency-path: studio/frontend/package-lock.json + # Run the structural lockfile scan BEFORE npm ci. A compromised + # tarball runs its `prepare` / `postinstall` during `npm ci`, + # so any catch has to fire upstream of that. The scanner is + # pure-Python read-only; safe to call ahead of every install. + - name: Lockfile supply-chain audit (pre-install scan) + working-directory: ${{ github.workspace }} + run: python3 scripts/lockfile_supply_chain_audit.py + - name: Lockfile must agree with package.json (npm ci is strict) run: npm ci --no-fund --no-audit diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index a1b54d6e65..19e1ab9cc5 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -94,6 +94,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache @@ -331,6 +333,8 @@ jobs: - name: Download GGUF if cache miss if: steps.cache-gguf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p gguf-cache @@ -637,6 +641,8 @@ jobs: - name: Prime HF_HOME with the GGUF + mmproj if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index 28a491840b..98596f374a 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -65,6 +65,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 066ddf87b8..97efe3e74d 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -88,6 +88,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache @@ -325,6 +327,8 @@ jobs: - name: Download GGUF if cache miss if: steps.cache-gguf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p gguf-cache @@ -679,13 +683,24 @@ jobs: - name: Prime HF_HOME with the GGUF + mmproj if: steps.cache-hf.outputs.cache-hit != 'true' + # Authenticated + parallel: shared macos-14 NAT egress stalls + # multi-GB anonymous downloads. + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$GGUF_FILE" + hf download "$GGUF_REPO" "$GGUF_FILE" & + MODEL_PID=$! HF_HUB_ENABLE_HF_TRANSFER=1 \ - hf download "$GGUF_REPO" "$MMPROJ_FILE" + hf download "$GGUF_REPO" "$MMPROJ_FILE" & + MMPROJ_PID=$! + wait "$MODEL_PID" + wait "$MMPROJ_PID" + # Fail loud on a partial download instead of in the next step. + find hf-cache -name "$GGUF_FILE" -o -name "$MMPROJ_FILE" \ + | xargs -I{} ls -lhL {} - name: Install Studio (--local, --no-torch) env: diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 75e958e023..c921ddf63e 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -65,6 +65,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml index d517a5f454..23b57d7e09 100644 --- a/.github/workflows/studio-tauri-smoke.yml +++ b/.github/workflows/studio-tauri-smoke.yml @@ -69,6 +69,9 @@ jobs: echo "$out" [ "$out" = "tauri-cli 2.10.1" ] || { echo "::error::expected tauri-cli 2.10.1, got $out"; exit 1; } + - name: Lockfile supply-chain audit (pre-install scan) + run: python3 scripts/lockfile_supply_chain_audit.py + - name: Frontend build (npm ci, vite) working-directory: studio/frontend run: | diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 6c4c66acd3..756eea64b2 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -79,6 +79,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index db2e8a26a0..d9ed5d5594 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -72,6 +72,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index e1406b7f45..b33b6f9563 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -82,6 +82,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache @@ -382,6 +384,8 @@ jobs: - name: Download GGUF if cache miss if: steps.cache-gguf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p gguf-cache @@ -776,6 +780,8 @@ jobs: - name: Prime HF_HOME with the GGUF + mmproj if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index c550f04827..a5a4753ba5 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -81,6 +81,8 @@ jobs: - name: Prime HF_HOME with the GGUF if: steps.cache-hf.outputs.cache-hit != 'true' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | python -m pip install --upgrade huggingface_hub hf_transfer mkdir -p hf-cache diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index ff3218bba0..b14a759916 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -203,8 +203,22 @@ jobs: with: { path: unsloth } - name: Clone unsloth-zoo @ main run: | - git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ - "$RUNNER_TEMP/unsloth-zoo" + # github.com occasionally 500s on the git fetch; retry so a + # single upstream blip does not fail CI. + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: python-version: '3.12' diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml index 983070ae13..dad8670393 100644 --- a/.github/workflows/wheel-smoke.yml +++ b/.github/workflows/wheel-smoke.yml @@ -53,6 +53,9 @@ jobs: with: python-version: '3.12' + - name: Lockfile supply-chain audit (pre-install scan) + run: python3 scripts/lockfile_supply_chain_audit.py + - name: Build frontend run: | cd studio/frontend diff --git a/.semgrep/unsloth-rules.yml b/.semgrep/unsloth-rules.yml deleted file mode 100644 index 654ff9a490..0000000000 --- a/.semgrep/unsloth-rules.yml +++ /dev/null @@ -1,183 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# -# Custom Semgrep rules for unsloth + studio backend. The off-the-shelf -# rule packs (p/python, p/javascript, p/supply-chain, p/security-audit) -# wired into the security-audit workflow already cover the common -# patterns. These rules add catches for the *specific* shape of recent -# CVEs in the broader Python ML / dev-tools stack -- so if we ever -# introduce a similar bug ourselves, CI lights up. -# -# Run locally: -# pip install 'semgrep>=1.95' -# semgrep --config .semgrep/unsloth-rules.yml studio/backend unsloth scripts -# -# Wired into CI via .github/workflows/security-audit.yml's Semgrep step. - -rules: - # ───────────────────────────────────────────────────────────────── - # langchain-core CVE-2025-68664 shape: - # `dumps()` / `dumpd()` over a user-controlled dict that may carry - # the `lc` marker key -> deserialization injection on the round - # trip. Catch any json.dumps / pickle.dumps / yaml.dump on data - # that flowed through a Request/WebSocket payload. - # ───────────────────────────────────────────────────────────────── - - id: unsloth-deserialize-roundtrip - message: >- - Serializing user-controlled data with langchain-style `dumps` - can re-instantiate arbitrary classes when deserialized. See - langchain-core CVE-2025-68664. Sanitize / strip `lc` marker keys - before dumping, or use a strict schema (Pydantic) instead. - severity: WARNING - languages: [python] - patterns: - - pattern-either: - - pattern: langchain_core.load.dumps($DATA, ...) - - pattern: langchain_core.load.dumpd($DATA, ...) - - pattern: dumps($DATA) - - pattern: dumpd($DATA) - - metavariable-pattern: - metavariable: $DATA - patterns: - - pattern-either: - - pattern: request.$F - - pattern: payload - - pattern: body - - pattern: data - - pattern: input - - # ───────────────────────────────────────────────────────────────── - # n8n CVE-2025-68668 shape: - # `_pyodide._base.eval_code(...)` or any private/underscore call - # into pyodide internals that escapes the public sandbox API. - # ───────────────────────────────────────────────────────────────── - - id: unsloth-pyodide-private-eval - message: >- - Calling `_pyodide._base.eval_code` (or any `_pyodide.`) - bypasses the public Pyodide sandbox -- this is how n8n - CVE-2025-68668 (CVSS 9.9) escaped the Code Node's blocklist. - Use the documented sandbox API (`pyodide.runPython`) and rely - on web-worker isolation for untrusted input. - severity: ERROR - languages: [python, javascript, typescript] - patterns: - - pattern-either: - - pattern: _pyodide._base.eval_code(...) - - pattern: $X._pyodide.$Y(...) - - # ───────────────────────────────────────────────────────────────── - # marimo CVE-2026-39987 shape: - # FastAPI / Starlette WebSocket route that accepts connections - # without checking auth -- in marimo this dropped a PTY shell to - # any unauthenticated attacker. - # ───────────────────────────────────────────────────────────────── - - id: unsloth-websocket-no-auth - message: >- - WebSocket route accepts connections without an auth check. - marimo CVE-2026-39987 was a pre-auth WebSocket on - `/terminal/ws` that handed a full PTY shell to any - unauthenticated peer. Add a Depends(get_current_user) / - `await websocket.headers.get("authorization")` gate before - `await websocket.accept()`. - severity: WARNING - languages: [python] - patterns: - - pattern: | - @$APP.websocket("...") - async def $F(websocket: WebSocket, ...): - ... - await websocket.accept() - ... - - pattern-not-inside: | - @$APP.websocket("...") - async def $F(websocket: WebSocket, ..., $USER = Depends(...)): - ... - - pattern-not-inside: | - @$APP.websocket("...") - async def $F(websocket: WebSocket, ...): - ... - if not $AUTH: - ... - await websocket.accept() - - # ───────────────────────────────────────────────────────────────── - # litellm 1.82.7 shape: - # `subprocess.Popen` of a child Python interpreter that reads - # stdin from a network response (the C2-fetch-then-exec dropper - # pattern). Catches both `Popen([sys.executable, ...], stdin=...)` - # and `Popen("python ...", stdin=...)` variants. - # ───────────────────────────────────────────────────────────────── - - id: unsloth-popen-network-stdin - message: >- - Spawning a Python interpreter that reads its program from a - network call is the canonical fetch-and-exec dropper (litellm - 1.82.7 used this exact shape). Almost never legitimate inside a - package's import path. - severity: ERROR - languages: [python] - pattern-either: - - pattern: | - subprocess.Popen([..., $PY, ...], stdin=$NET, ...) - - pattern: | - subprocess.run([..., $PY, ...], input=$NET, ...) - - # ───────────────────────────────────────────────────────────────── - # Shai-Hulud / ForceMemo shape: - # programmatic write of a `.github/workflows/*.yml` file from - # inside our own Python source. We never write workflows - # programmatically; if a contributor ever does, they're probably - # re-implementing the worm pattern. - # ───────────────────────────────────────────────────────────────── - - id: unsloth-write-github-workflow - message: >- - Code that programmatically writes into `.github/workflows/` - from within unsloth itself is the Shai-Hulud / ForceMemo - self-propagation pattern. If you legitimately need a workflow - template, ship it under examples/ or templates/ instead. - severity: ERROR - languages: [python] - patterns: - - pattern-either: - - pattern: open("$P", ...) - - pattern: Path("$P").write_text(...) - - pattern: open("$P", "w", ...) - - metavariable-regex: - metavariable: $P - regex: \.github/workflows/.*\.ya?ml - - # ───────────────────────────────────────────────────────────────── - # Pickle-from-network shape: classic deserialization sink that - # several recent ML pipeline CVEs hit (mlflow, pyzmq, ray serve). - # ───────────────────────────────────────────────────────────────── - - id: unsloth-pickle-from-network - message: >- - `pickle.loads` on bytes that flowed from a network response is - arbitrary code execution. Use `safetensors` or a strict - schema (Pydantic / msgspec) instead. ML frameworks have shipped - multiple CVEs of this exact shape (mlflow, ray serve, pyzmq). - severity: ERROR - languages: [python] - pattern-either: - - pattern: pickle.loads($X.content) - - pattern: pickle.loads($X.text.encode(...)) - - pattern: pickle.loads(requests.get(...).content) - - pattern: pickle.load(urllib.request.urlopen(...)) - - # ───────────────────────────────────────────────────────────────── - # Subprocess shell=True with f-string / format / concat -- command - # injection if any interpolated value comes from user input. - # ───────────────────────────────────────────────────────────────── - - id: unsloth-shell-true-interpolation - message: >- - `subprocess` call with `shell=True` and an interpolated command - string is command injection if any input is user-controlled. - Pass argv list instead, or use shlex.quote on each part. - severity: WARNING - languages: [python] - pattern-either: - - pattern: subprocess.run(f"...", shell=True, ...) - - pattern: subprocess.Popen(f"...", shell=True, ...) - - pattern: subprocess.call(f"...", shell=True, ...) - - pattern: os.system(f"...") - - pattern: subprocess.run("..." + $X, shell=True, ...) - - pattern: subprocess.run("...{}...".format(...), shell=True, ...) diff --git a/build.sh b/build.sh index cf8aa02910..1558dca240 100644 --- a/build.sh +++ b/build.sh @@ -2,6 +2,10 @@ set -euo pipefail +# PyPI/Studio release publishing must use `./build.sh publish` (or an +# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Studio +# artifacts include the display-only Studio release version. + # 1. Build frontend (Vite outputs to dist/) cd studio/frontend @@ -70,10 +74,33 @@ cd ../.. # 2. Clean old artifacts rm -rf build dist *.egg-info -# 3. Build wheel +# 3. Stamp display-only Studio release metadata for packaged builds. +_STUDIO_BUILD_INFO="studio/backend/utils/_studio_release_build.py" +_STUDIO_BUILD_INFO_BACKUP="$(mktemp)" +cp "$_STUDIO_BUILD_INFO" "$_STUDIO_BUILD_INFO_BACKUP" +_restore_studio_build_info() { + cp "$_STUDIO_BUILD_INFO_BACKUP" "$_STUDIO_BUILD_INFO" 2>/dev/null || true + rm -f "$_STUDIO_BUILD_INFO_BACKUP" +} +trap _restore_studio_build_info EXIT + +if [ "${1:-}" = "publish" ]; then + STUDIO_STAMPED_VERSION="$(python scripts/stamp_studio_release.py --require-release)" +else + STUDIO_STAMPED_VERSION="$(python scripts/stamp_studio_release.py)" +fi + +# 4. Build wheel/sdist python -m build -# 4. Optionally publish +if [ "${1:-}" = "publish" ]; then + python scripts/stamp_studio_release.py --verify-dist dist --expected "$STUDIO_STAMPED_VERSION" +fi + +_restore_studio_build_info +trap - EXIT + +# 5. Optionally publish if [ "${1:-}" = "publish" ]; then python -m twine upload dist/* fi diff --git a/scripts/lockfile_supply_chain_audit.py b/scripts/lockfile_supply_chain_audit.py new file mode 100755 index 0000000000..e52183214e --- /dev/null +++ b/scripts/lockfile_supply_chain_audit.py @@ -0,0 +1,486 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Lockfile supply-chain audit for the Studio frontend and Tauri shell. + +Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a +lockfile contains patterns that indicate the kind of supply-chain +injection seen in the npm Shai-Hulud waves and the cargo +crates.io brand-squat attempts. + +What it checks +============== + +studio/frontend/package-lock.json (lockfileVersion 2 or 3): + + 1. `resolved` URL origin. Every entry must resolve through + `https://registry.npmjs.org/`. Direct GitHub-hosted dependencies + (`git+ssh://`, `git+https://`, `github:owner/repo#sha`, + `file:`, `http://`) are refused -- npm's TanStack incident used + exactly this vector to land an unaudited GitHub commit hash as + an optional dependency. + + 2. `integrity` field presence. Every non-workspace entry must carry + an `integrity` SHA. A missing integrity means the registry can + swap the tarball after lockfile generation and CI will not + notice. + + 3. Known IOC strings. A hardcoded set of indicator-of-compromise + substrings is grepped across the entire lockfile body (file + names, dependency keys, URLs). The list is updated as new + campaigns surface. Catching one means the local install was + about to pull a publicly-known malicious release. + +studio/src-tauri/Cargo.lock: + + 4. `source` field origin. Every entry with a `source` must point at + `registry+https://github.com/rust-lang/crates.io-index`. Direct + git sources (`git+https://...`) and `path+...` for cross-crate + paths warrant manual review and are flagged. + + 5. Known cargo IOC strings. Same idea as (3), separate list. + +Exit codes +========== + + 0 no findings, or an opt-out env var (UNSLOTH_LOCKFILE_AUDIT_SKIP=1) + is set + 1 one or more findings; stderr lists them with file path and line + number where derivable + 2 internal error (missing dependency, malformed JSON, etc.) + +Operational stance +================== + +This scanner only PARSES the lockfiles -- it never executes anything +in them, never resolves anything against the network. Safe to run +ahead of every `npm ci`. The IOC list is short by design; this +complements (not replaces) `npm audit`, OSV-Scanner, and the +advisory-DB pipeline in `.github/workflows/security-audit.yml`. The +shape of the catch is "we refuse to proceed because the lockfile +itself is shaped wrong", which fires before any third-party install +script gets a chance to run on the runner. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +# ───────────────────────────────────────────────────────────────────── +# Known IOC strings (case-sensitive substring match). +# ───────────────────────────────────────────────────────────────────── +# +# Keep these short and FACTUAL. Each entry is tied to a public advisory +# and is the literal string an attacker would have to embed for the +# attack to work. Adding speculative or generic patterns here would +# generate false positives on dependency upgrades. +NPM_IOC_STRINGS: tuple[str, ...] = ( + # Shai-Hulud TanStack wave -- May 11, 2026 (GHSA-g7cv-rxg3-hmpx). + "router_init.js", + "tanstack_runner.js", + "router_runtime.js", + "@tanstack/setup", + "github:tanstack/router#79ac49eedf774dd4b0cfa308722bc463cfe5885c", + # Exfiltration endpoints observed across both Shai-Hulud waves. + "filev2.getsession.org", + "getsession.org/file/", + # Campaign markers; the worm tarballs print this to stdout on run. + "A Mini Shai-Hulud has Appeared", +) + +CARGO_IOC_STRINGS: tuple[str, ...] = ( + # Reserved for future cargo-side incidents. Empty by default -- + # `source` origin check below catches the structural pattern. +) + + +# ───────────────────────────────────────────────────────────────────── +# Allowed lockfile origins. +# ───────────────────────────────────────────────────────────────────── +NPM_REGISTRY_PREFIX = "https://registry.npmjs.org/" + +# Tarballs are also fetched from this mirror on some GH Actions cached +# runs (npm rewrites the resolved URL on cache hit). Allow either. +NPM_REGISTRY_PREFIXES_ALLOWED: tuple[str, ...] = (NPM_REGISTRY_PREFIX,) + +CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index" + + +# ───────────────────────────────────────────────────────────────────── +# Cargo non-registry source allowlist. +# ───────────────────────────────────────────────────────────────────── +# +# Each entry is `(crate_name, exact_source_string)`. The crate must +# match by name AND the source must match the full pinned-SHA string +# verbatim. Bumping the commit SHA forces a re-review here: the +# scanner fires until the new SHA is appended. +# +# Studio's Tauri shell pulls `fix-path-env` directly from +# tauri-apps/fix-path-env-rs because the crate is not published to +# crates.io. The pinned commit (c4c45d5) was reviewed at the time it +# landed; future bumps need explicit approval. +CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = ( + ( + "fix-path-env", + "git+https://github.com/tauri-apps/fix-path-env-rs#" + "c4c45d503ea115a839aae718d02f79e7c7f0f673", + ), +) + + +# ───────────────────────────────────────────────────────────────────── +# Finding container. +# ───────────────────────────────────────────────────────────────────── + + +class Finding: + __slots__ = ("path", "package", "kind", "detail") + + def __init__(self, path: str, package: str, kind: str, detail: str) -> None: + self.path = path + self.package = package + self.kind = kind + self.detail = detail + + def __str__(self) -> str: + return ( + f" [{self.kind}] {self.path}\n" + f" package: {self.package}\n" + f" detail: {self.detail}" + ) + + +# ───────────────────────────────────────────────────────────────────── +# package-lock.json audit. +# ───────────────────────────────────────────────────────────────────── + + +def audit_npm_lockfile(path: Path) -> list[Finding]: + findings: list[Finding] = [] + if not path.exists(): + return findings + + raw = path.read_text(encoding = "utf-8") + try: + lock = json.loads(raw) + except json.JSONDecodeError as exc: + findings.append( + Finding( + path = str(path), + package = "", + kind = "malformed-lockfile", + detail = f"could not parse as JSON: {exc}", + ) + ) + return findings + + lockfile_version = lock.get("lockfileVersion") + if lockfile_version not in (2, 3): + findings.append( + Finding( + path = str(path), + package = "", + kind = "unsupported-lockfile-version", + detail = (f"only lockfileVersion 2 or 3 audited; got {lockfile_version}"), + ) + ) + + packages = lock.get("packages") or {} + for key, entry in packages.items(): + # The empty key "" is the project root; workspace entries use + # keys like "node_modules/foo" or "studio/frontend/sub-pkg". + # Skip the project root (it has no `resolved`). + if key == "": + continue + if entry.get("link"): + # Workspace symlink; no tarball to resolve. + continue + + resolved = entry.get("resolved") + # Entries living inside another package's `node_modules/` + # tree are bundled fold-ins -- the parent's tarball ships + # their source verbatim and the parent's `integrity` covers + # the whole subtree. npm represents them in lockfileVersion 3 + # as nested entries with no `resolved` and no `integrity` of + # their own. Treat them as transparent to this audit. + nested = key.count("/node_modules/") >= 1 + + # 1. resolved-URL origin. + if resolved is None: + if nested or entry.get("bundled"): + # Bundled / fold-in entry; covered by parent integrity. + pass + elif entry.get("version"): + # Top-level entry without a resolved URL is suspicious. + findings.append( + Finding( + path = str(path), + package = key, + kind = "missing-resolved-url", + detail = ( + f"version={entry['version']!r} but no `resolved` " + "field; lockfile is incomplete" + ), + ) + ) + else: + if not any(resolved.startswith(p) for p in NPM_REGISTRY_PREFIXES_ALLOWED): + findings.append( + Finding( + path = str(path), + package = key, + kind = "non-registry-resolved-url", + detail = ( + f"resolved={resolved!r}; only " + f"{NPM_REGISTRY_PREFIX} is permitted. Direct " + "GitHub / git / file references are the " + "Shai-Hulud injection vector." + ), + ) + ) + + # 2. integrity-hash presence. + if resolved is not None and not entry.get("integrity"): + findings.append( + Finding( + path = str(path), + package = key, + kind = "missing-integrity-hash", + detail = ( + "no `integrity` field; npm cannot verify the " + "tarball SHA against the registry-published hash" + ), + ) + ) + + # 3. Known IOC strings: scan the raw file body so we hit fields the + # structural pass above doesn't enumerate (scripts, optional + # dependencies, etc.). Cheap and complete. + for ioc in NPM_IOC_STRINGS: + if ioc in raw: + # Best-effort line number lookup. + line_no = _first_line_containing(raw, ioc) + findings.append( + Finding( + path = f"{path}:{line_no}" if line_no else str(path), + package = "", + kind = "known-ioc-string", + detail = ( + f"matched known IOC substring {ioc!r}; this is " + "a public indicator of a recent supply-chain " + "compromise. Refuse to install." + ), + ) + ) + + return findings + + +def _first_line_containing(text: str, needle: str) -> int | None: + for i, line in enumerate(text.splitlines(), start = 1): + if needle in line: + return i + return None + + +# ───────────────────────────────────────────────────────────────────── +# Cargo.lock audit. +# ───────────────────────────────────────────────────────────────────── + + +# Cargo.lock is TOML; parse with stdlib tomllib (Python 3.11+). The +# studio's Tauri shell already requires a modern toolchain so this is +# always available where CI runs. +_PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$") + + +def audit_cargo_lockfile(path: Path) -> list[Finding]: + findings: list[Finding] = [] + if not path.exists(): + return findings + + raw = path.read_text(encoding = "utf-8") + try: + import tomllib # type: ignore[import-not-found] + except ImportError: + # Python <3.11; fall back to a tomli shim if importable. + try: + import tomli as tomllib # type: ignore[no-redef] + except ImportError: + findings.append( + Finding( + path = str(path), + package = "", + kind = "missing-toml-parser", + detail = ( + "Python 3.11+ tomllib or tomli is required to " + "parse Cargo.lock; install tomli or upgrade " + "Python before re-running this audit" + ), + ) + ) + return findings + + try: + lock = tomllib.loads(raw) + except Exception as exc: + findings.append( + Finding( + path = str(path), + package = "", + kind = "malformed-lockfile", + detail = f"could not parse as TOML: {exc}", + ) + ) + return findings + + for entry in lock.get("package", []): + name = entry.get("name") or "" + version = entry.get("version") or "" + source = entry.get("source") + # Workspace-local crates have no `source` field; skip them. + if source is None: + continue + if source != CARGO_REGISTRY_SOURCE: + if (name, source) in CARGO_SOURCE_ALLOWLIST: + # Pre-approved non-registry source pinned by SHA. + pass + else: + findings.append( + Finding( + path = str(path), + package = f"{name}@{version}", + kind = "non-registry-cargo-source", + detail = ( + f"source={source!r}; only " + f"{CARGO_REGISTRY_SOURCE!r} is permitted " + "by default, and no allowlist entry covers " + "this crate. If the source is legitimate, " + "add `(name, source)` to " + "CARGO_SOURCE_ALLOWLIST after reviewing the " + "pinned commit." + ), + ) + ) + if not entry.get("checksum") and source == CARGO_REGISTRY_SOURCE: + findings.append( + Finding( + path = str(path), + package = f"{name}@{version}", + kind = "missing-cargo-checksum", + detail = ( + "registry crate without checksum; cargo cannot " + "verify the downloaded source against the " + "registry-published SHA" + ), + ) + ) + + for ioc in CARGO_IOC_STRINGS: + if ioc in raw: + line_no = _first_line_containing(raw, ioc) + findings.append( + Finding( + path = f"{path}:{line_no}" if line_no else str(path), + package = "", + kind = "known-ioc-string", + detail = f"matched known IOC substring {ioc!r}", + ) + ) + + return findings + + +# ───────────────────────────────────────────────────────────────────── +# CLI. +# ───────────────────────────────────────────────────────────────────── + + +DEFAULT_NPM_LOCKFILES = ("studio/frontend/package-lock.json",) +DEFAULT_CARGO_LOCKFILES = ("studio/src-tauri/Cargo.lock",) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description = "Pre-install lockfile supply-chain audit.", + ) + parser.add_argument( + "--root", + default = str(REPO_ROOT), + help = "Repo root (default: parent of this script).", + ) + parser.add_argument( + "--npm-lockfile", + action = "append", + default = None, + help = ( + "Path to a package-lock.json (repeatable). " + "Default: studio/frontend/package-lock.json." + ), + ) + parser.add_argument( + "--cargo-lockfile", + action = "append", + default = None, + help = ( + "Path to a Cargo.lock (repeatable). " + "Default: studio/src-tauri/Cargo.lock." + ), + ) + args = parser.parse_args(argv) + + if os.environ.get("UNSLOTH_LOCKFILE_AUDIT_SKIP") == "1": + print( + "[lockfile-audit] UNSLOTH_LOCKFILE_AUDIT_SKIP=1; " + "audit skipped (expected only for local triage)", + flush = True, + ) + return 0 + + root = Path(args.root).resolve() + npm_paths = [root / p for p in (args.npm_lockfile or DEFAULT_NPM_LOCKFILES)] + cargo_paths = [root / p for p in (args.cargo_lockfile or DEFAULT_CARGO_LOCKFILES)] + + all_findings: list[Finding] = [] + for p in npm_paths: + print(f"[lockfile-audit] npm: {p}", flush = True) + all_findings.extend(audit_npm_lockfile(p)) + for p in cargo_paths: + print(f"[lockfile-audit] cargo: {p}", flush = True) + all_findings.extend(audit_cargo_lockfile(p)) + + if not all_findings: + print( + f"[lockfile-audit] OK: 0 findings across " + f"{len(npm_paths)} npm + {len(cargo_paths)} cargo lockfile(s)", + flush = True, + ) + return 0 + + print( + f"\n[lockfile-audit] FAIL: {len(all_findings)} finding(s):\n", + file = sys.stderr, + ) + for f in all_findings: + print(str(f), file = sys.stderr) + print(file = sys.stderr) + print( + "[lockfile-audit] Refusing to proceed. Each finding above is " + "either a structural lockfile anomaly or a public indicator-of-" + "compromise. Investigate before running `npm ci` or `cargo fetch`.", + file = sys.stderr, + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py new file mode 100755 index 0000000000..97b5ffa9a4 --- /dev/null +++ b/scripts/scan_npm_packages.py @@ -0,0 +1,1201 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# .github/workflows/security-audit.yml's npm-scan-packages job depends +# on this file existing at scripts/scan_npm_packages.py. + +"""scan_npm_packages.py -- npm-side content scanner. + +Counterpart to scripts/scan_packages.py for the pip ecosystem. Reads +studio/frontend/package-lock.json, downloads each resolved tarball +DIRECTLY from registry.npmjs.org (never via `npm install` -- no +lifecycle scripts ever run), verifies the lockfile integrity hash, +unpacks each tarball into a sandboxed temp dir behind size / count / +path-escape / symlink guards, and pattern-scans the extracted file +contents for the signatures common to npm supply-chain attacks: + + - Lifecycle (preinstall / install / postinstall / prepare) scripts + in any package.json that fetch + execute external code. + - C2 / exfiltration hosts (getsession.org, AWS IMDS endpoints, + Kubernetes ServiceAccount token paths, GitHub Actions OIDC, + HashiCorp Vault endpoints). + - Credential-stealing references (~/.npmrc, ~/.aws/credentials, + GITHUB_TOKEN / NPM_TOKEN in JS sources). + - Known IOC filenames from public advisories + (router_init.js, tanstack_runner.js, router_runtime.js). + - Obfuscation shapes (large single JS in package root with a low + whitespace ratio + Function/eval against a base64-decoded blob). + +Safety stance +============= + +This script ingests attacker-controlled archives. Every parse path +assumes the worst: + + 1. Downloads ONLY from `registry.npmjs.org`. Any tarball URL with a + different hostname is refused without fetching. + 2. Tarball download is size-capped (HARD_MAX_TARBALL_BYTES default + 64 MiB). HEAD-style probe via the Content-Length response header + plus a chunked read that aborts on overflow. + 3. SHA-512 integrity verified against the lockfile entry BEFORE the + tarball is even opened. A mismatch aborts that package -- the + scanner does not "fall back" to the registry-published hash. + 4. tar extraction goes through `safe_extract`: + - rejects symbolic links (`SYMTYPE`, `LNKTYPE`) + - rejects absolute paths, `..` traversal, paths outside the + extract root after resolution + - rejects character / block / FIFO devices + - per-file uncompressed size cap (HARD_MAX_FILE_BYTES, default + 8 MiB) AND cumulative cap (HARD_MAX_TOTAL_BYTES, default + 128 MiB) AND member-count cap (HARD_MAX_MEMBERS, default + 50_000) + - tar reads happen via `tarfile.open(mode='r|gz')` streaming + so an oversized file is detected before write + 5. NOTHING from the extracted tree is ever executed. Files are read + as raw bytes, decoded with `errors='replace'`, and grepped. We + never call `node`, `eval`, `compile`, `subprocess.run`, + `os.system`, or anything that would touch the tarball's + declared scripts. + 6. Tempdir is created with `tempfile.mkdtemp(prefix='npm-scan-')`, + fully resolved with .resolve(), and registered with atexit to be + wiped on every termination path. + 7. Stdlib only. No third-party deps -- adding one would itself be a + supply-chain liability. + +Exit codes +========== + + 0 no findings of severity HIGH or higher + 1 one or more HIGH/CRITICAL findings (or pre-scan structural + anomalies -- non-registry resolved URL, missing integrity) + 2 internal error (lockfile missing, integrity mismatch on + download, malformed tarball, etc.) + +The script is meant to be run in CI on every PR that touches +package-lock.json and on a nightly schedule. +""" + +from __future__ import annotations + +import argparse +import atexit +import base64 as _b64 # imported only so the IOC string-scan can detect it +import hashlib +import io +import json +import os +import re +import shutil +import sys +import tarfile +import tempfile +import urllib.parse +import urllib.request +from dataclasses import dataclass, field +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# ───────────────────────────────────────────────────────────────────── +# Hard caps (deliberately conservative; npm tarballs in this repo are +# all well under these limits, so a packaging spike is noticeable). +# ───────────────────────────────────────────────────────────────────── +# Caps calibrated against the real Studio frontend transitive closure: +# - typescript.js is 9.1 MB (TS compiler bundled into one file) +# - mermaid 11.x dist/mermaid.js.map is ~12 MB (sourcemap) +# - lightningcss-linux-x64-{gnu,musl}.node is 10 MB +# - rolldown bindings (.node) are 18-26 MB per platform +# - @next/swc-*.node is ~137 MB (rust-compiled SWC engine) +# - next.js cumulative bundle is ~134 MB (turbopack compiled) +# +# Native binaries (.node, .wasm, .so, .dll, .dylib) are GENUINELY +# huge and not amenable to text pattern scanning -- we extract them +# only to verify the tarball integrity over the full archive, then +# skip them in scan_extracted_tree. They get a much higher per-file +# cap. Text files (JS/TS/JSON/etc) keep the tight cap because the +# pattern scanner runs over them and a 9.1 MB typescript.js is the +# legitimate ceiling. +HARD_MAX_TARBALL_BYTES = 256 * 1024 * 1024 # 256 MiB compressed +HARD_MAX_TEXT_FILE_BYTES = 16 * 1024 * 1024 # 16 MiB per text file +HARD_MAX_BINARY_FILE_BYTES = 256 * 1024 * 1024 # 256 MiB per .node etc +HARD_MAX_TOTAL_BYTES = 512 * 1024 * 1024 # 512 MiB cumulative +HARD_MAX_MEMBERS = 50_000 # entries per tarball +HARD_HTTP_TIMEOUT_S = 60 # per request + +# Native-binary / compiled-asset suffixes that bypass the text cap. +# This is the SUFFIX shortlist; the content-magic check below covers +# extensionless executables (biome) and versioned shared libraries +# (libvips-cpp.so.8.17.3) that the suffix list misses. +_BINARY_SUFFIXES = ( + ".node", + ".wasm", + ".so", + ".dll", + ".dylib", + ".exe", + ".a", + ".lib", + ".o", + ".obj", + ".bin", + ".dat", + ".woff", + ".woff2", + ".ttf", + ".otf", + ".eot", + ".png", + ".jpg", + ".jpeg", + ".gif", + ".webp", + ".ico", + ".mp3", + ".mp4", + ".webm", + ".zip", + ".tar", + ".gz", + ".tgz", + ".xz", + ".bz2", +) + +# Versioned shared libraries: libfoo.so.1.2.3 / libfoo.dylib.1.2. +_VERSIONED_LIB = re.compile( + r"\.(?:so|dylib)(?:\.\d+)+$", + re.IGNORECASE, +) + +# Magic numbers at offset 0 that identify common executable formats. +# We sniff the first ~16 bytes of every member to catch extensionless +# binaries (eg `package/biome`, `package/bin/foo`). +_BINARY_MAGICS = ( + b"\x7fELF", # ELF (Linux executable / .so) + b"MZ", # PE / .exe / .dll (DOS header prefix) + b"\xfe\xed\xfa\xce", # Mach-O 32 BE + b"\xfe\xed\xfa\xcf", # Mach-O 64 BE + b"\xce\xfa\xed\xfe", # Mach-O 32 LE + b"\xcf\xfa\xed\xfe", # Mach-O 64 LE + b"\xca\xfe\xba\xbe", # Mach-O fat / Java class (also starts with this) + b"\x00asm", # WASM + b"PK\x03\x04", # ZIP / JAR / nupkg / xpi + b"PK\x05\x06", # ZIP (empty) + b"\x1f\x8b", # gzip + b"BZh", # bzip2 + b"\xfd7zXZ", # xz + b"7z\xbc\xaf\x27\x1c", # 7zip + b"\x89PNG", # PNG + b"\xff\xd8\xff", # JPEG + b"GIF8", # GIF + b"RIFF", # WAV / WEBP / AVI container + b"\x00\x00\x01\x00", # ICO + b"OggS", # Ogg + b"\x1aE\xdf\xa3", # Matroska / WebM +) + + +def _looks_binary(name: str, header: bytes) -> bool: + """True if `name` or first bytes suggest a non-text file.""" + lower = name.lower() + if lower.endswith(_BINARY_SUFFIXES): + return True + if _VERSIONED_LIB.search(lower): + return True + for magic in _BINARY_MAGICS: + if header.startswith(magic): + return True + # Null-byte density: real text files almost never carry NULs. + if header and (header.count(b"\x00") / len(header)) > 0.02: + return True + return False + + +ALLOWED_DOWNLOAD_HOST = "registry.npmjs.org" + +# ───────────────────────────────────────────────────────────────────── +# Severities + finding shape (mirrors scripts/scan_packages.py). +# ───────────────────────────────────────────────────────────────────── +CRITICAL = "CRITICAL" +HIGH = "HIGH" +MEDIUM = "MEDIUM" +INFO = "INFO" +_SEVERITY_RANK = {CRITICAL: 0, HIGH: 1, MEDIUM: 2, INFO: 3} + + +@dataclass +class Finding: + severity: str + package: str # name@version + filename: str # relative path inside the tarball + pattern: str # what matched + evidence: str = "" # short surrounding snippet + detail: str = "" # human-readable description + + def __str__(self) -> str: + head = f" [{self.severity}] {self.package} :: {self.filename}" + body = f" pattern: {self.pattern}" + if self.detail: + body += f"\n detail: {self.detail}" + if self.evidence: + ev = self.evidence + if len(ev) > 240: + ev = ev[:240] + "..." + body += f"\n evidence: {ev!r}" + return f"{head}\n{body}" + + +@dataclass +class PackageEntry: + name: str + version: str + resolved: str + integrity: str | None + lockfile_key: str + + @property + def display(self) -> str: + return f"{self.name}@{self.version}" + + +# ───────────────────────────────────────────────────────────────────── +# IOC patterns. Two flavours: +# - HOSTS / TOKEN_PATHS: high-confidence substrings; near-zero FP rate +# - JS_PATTERNS / SCRIPT_PATTERNS: regex; tuned to recent campaigns +# Keep this list short and factual. Speculative patterns spam the +# false-positive ledger and dull the signal. +# ───────────────────────────────────────────────────────────────────── + + +# Substring (case-sensitive) -> (severity, detail). +KNOWN_IOC_STRINGS: dict[str, tuple[str, str]] = { + # Shai-Hulud TanStack wave (2026-05-11, GHSA-g7cv-rxg3-hmpx). + "router_init.js": (HIGH, "filename associated with TanStack worm"), + "tanstack_runner.js": (HIGH, "filename associated with TanStack worm"), + "router_runtime.js": (HIGH, "filename associated with TanStack worm"), + "A Mini Shai-Hulud has Appeared": ( + CRITICAL, + "TanStack worm campaign stdout marker", + ), + "github:tanstack/router#79ac49eedf774dd4b0cfa308722bc463cfe5885c": ( + CRITICAL, + "TanStack worm dropper pinned commit", + ), + # Exfil hosts observed across both Shai-Hulud waves. + "filev2.getsession.org": (CRITICAL, "exfiltration C2 host"), + "getsession.org/file/": (CRITICAL, "exfiltration C2 endpoint"), +} + +# Cloud / k8s / CI credential surfaces. A bare substring match here +# false-positives on DEFENSIVE code -- e.g. langchain ships an SSRF +# protection module with a literal blocklist of IMDS IPs. We split +# these into two tiers: +# +# ALWAYS_BAD: substrings with no legitimate use anywhere in a +# dependency. A bare match is enough. +# +# NEEDS_CONTEXT: hosts/paths that DO appear legitimately in +# defensive code. We only fire when they co-occur with a fetch +# verb or appear inside an http URL -- that is the structural +# difference between "blocked address constant" and "exfil +# target". +# +# The dispatch lives in `_scan_cred_surface` below. + +CRED_HOST_ALWAYS_BAD: tuple[tuple[str, str], ...] = ( + ("registry.npmjs.org/-/npm/v1/tokens", "npm publish-token enumeration endpoint"), + ("ACTIONS_ID_TOKEN_REQUEST_URL", "GitHub Actions OIDC token-exchange endpoint env"), + ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "GitHub Actions OIDC token-exchange token env"), +) + +# Hosts that need fetch-verb or URL-scheme context to be malicious. +CRED_HOST_NEEDS_CONTEXT: tuple[tuple[str, str], ...] = ( + ("169.254.169.254", "AWS / GCP / Azure instance metadata service (IMDS)"), + ("169.254.170.2", "ECS task metadata service"), + ("metadata.google.internal", "GCE metadata service"), + ("vault.svc.cluster.local", "in-cluster HashiCorp Vault endpoint"), + ( + "/var/run/secrets/kubernetes.io/serviceaccount", + "Kubernetes ServiceAccount token path", + ), +) + +# Credentials a frontend package should NEVER need to read. Bare +# substring match is too noisy (object-treeify ships a `docker` dev +# script that mounts ~/.npmrc -- legitimate dev tooling, never run +# at install time). We instead surface these only when they appear +# inside a LIFECYCLE script (preinstall / install / postinstall / +# prepare), which is the only path that runs automatically on +# `npm ci`. See `scan_package_json` below. +CRED_PATH_SUBSTRINGS: tuple[tuple[str, str], ...] = ( + ("/.npmrc", "npm credentials file"), + ("/.aws/credentials", "AWS shared credentials file"), + ("/.ssh/id_rsa", "SSH private key"), + ("/.ssh/id_ed25519", "SSH private key"), + ("/.docker/config.json", "Docker registry credentials"), + ("/.kube/config", "Kubernetes kubeconfig"), +) + +# Fetch verbs whose presence near a metadata host upgrades a bare +# substring hit into an actionable finding. +_FETCH_VERBS_PAT = ( + r"(?:fetch|axios|XMLHttpRequest|got\b|undici|" + r"http\.get|https\.get|http\.request|https\.request|" + r"new\s+URL|url\.parse|net\.connect|" + r"\.request\s*\(|\.get\s*\(\s*['\"]\s*https?://)" +) + +# JS regex patterns (compile lazily). +_JS_FETCH_EVAL = re.compile( + r"""(?xs) + (?: + Function\s*\(\s*['"`] # new Function("...") + | eval\s*\(\s*['"`] + | \(\s*0\s*,\s*eval\s*\)\s*\( + ) + .{0,200} + (?:atob\s*\(|Buffer\s*\.from\s*\([^)]+,\s*['"]base64) + """, +) + +# `process.env.GITHUB_TOKEN` / `NPM_TOKEN` / `AWS_*` access in +# top-level / install-time code is suspicious. We also catch +# `os.environ["GITHUB_TOKEN"]` for the rare Python-in-npm postinstall. +_JS_ENV_TOKEN = re.compile( + r"""(process\.env\.|os\.environ\[?['"])(?: + GITHUB_TOKEN | GH_TOKEN | NPM_TOKEN | NODE_AUTH_TOKEN + | AWS_ACCESS_KEY_ID | AWS_SECRET_ACCESS_KEY | AWS_SESSION_TOKEN + | GOOGLE_APPLICATION_CREDENTIALS + | DOCKER_AUTH_CONFIG | VAULT_TOKEN + )['"]?\]?""", + re.VERBOSE, +) + +# Suspicious lifecycle-script payloads. Anything in a package.json +# `scripts` field that wgets/curls an external resource and executes +# it. We do NOT block ALL curl/wget in scripts (some legit packages +# fetch test fixtures into devDependencies), but we DO block the +# fetch+exec chain. +_LIFECYCLE_FETCH_EXEC = re.compile( + r"""(?xs) + (?:curl|wget|fetch|http\.get|axios\.get)\s+ # fetch verb + .{0,200} + (?:\|\s*(?:sh|bash|node|python|eval)\b # pipe to interpreter + | \&\&\s*(?:sh|bash|node|python|eval)\b # &&-chain to interpreter + | -o\s+\S+\s*&&\s*(?:sh|bash|node|python) # download then run + | --post-file\s+ + | \$\(.*\) # command-sub of fetched content + ) + """, +) + +# Obfuscation: large JS file that is mostly one line of base64-ish +# blob with a Function() / eval() bookend. Tuned against the +# router_init.js shape (2.3 MB obfuscated single-blob). +_OBFUSC_BLOB = re.compile( + r"""(?xs) + (?:Function|eval)\s*\(\s*['"`]? + [A-Za-z0-9+/=_-]{2048,} # >=2 KiB of b64-ish + """, +) + + +# ───────────────────────────────────────────────────────────────────── +# Lockfile parsing. +# ───────────────────────────────────────────────────────────────────── + + +def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]: + """Return (entries, structural_findings). + + Structural findings here are HIGH-severity refusals that should + short-circuit the scan -- a lockfile with non-registry resolved + URLs is itself a finding (covered by scripts/lockfile_supply_chain + _audit.py in detail; we surface a summary here so this scanner is + standalone-runnable). + """ + entries: list[PackageEntry] = [] + findings: list[Finding] = [] + + try: + lock = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, json.JSONDecodeError) as exc: + findings.append( + Finding( + severity = CRITICAL, + package = "", + filename = str(path), + pattern = "lockfile-unreadable", + detail = f"could not parse: {exc}", + ) + ) + return entries, findings + + if lock.get("lockfileVersion") not in (2, 3): + findings.append( + Finding( + severity = HIGH, + package = "", + filename = str(path), + pattern = "unsupported-lockfile-version", + detail = ( + f"only lockfileVersion 2 or 3 supported; got " + f"{lock.get('lockfileVersion')!r}" + ), + ) + ) + return entries, findings + + for key, entry in (lock.get("packages") or {}).items(): + if key == "" or entry.get("link"): + continue + # Nested fold-ins (deps inside another package's node_modules/) + # are covered by the parent tarball's integrity. Skip. + if key.count("/node_modules/") >= 1: + continue + resolved = entry.get("resolved") + if not resolved: + continue + # Strict registry origin check. lockfile_supply_chain_audit + # already catches this; double-defend here so this scanner + # cannot be tricked into fetching from an attacker-chosen URL. + parsed = urllib.parse.urlparse(resolved) + if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST: + findings.append( + Finding( + severity = CRITICAL, + package = key, + filename = str(path), + pattern = "non-registry-resolved-url", + detail = ( + f"resolved={resolved!r}; only " + f"https://{ALLOWED_DOWNLOAD_HOST}/ is " + "permitted. Refusing to download." + ), + ) + ) + continue + integrity = entry.get("integrity") + if not integrity: + findings.append( + Finding( + severity = HIGH, + package = key, + filename = str(path), + pattern = "missing-integrity-hash", + detail = "no `integrity` field; cannot verify download", + ) + ) + continue + # node_modules/@scope/name -> @scope/name; node_modules/name -> name + nm = "node_modules/" + name = key[len(nm) :] if key.startswith(nm) else key + version = entry.get("version") or "" + entries.append( + PackageEntry( + name = name, + version = version, + resolved = resolved, + integrity = integrity, + lockfile_key = key, + ) + ) + return entries, findings + + +# ───────────────────────────────────────────────────────────────────── +# Tarball download (registry-only, size-capped, integrity-verified). +# ───────────────────────────────────────────────────────────────────── + + +def _decode_integrity(integrity: str) -> tuple[str, bytes] | None: + """Parse SRI integrity 'sha512-' -> (algo, digest_bytes).""" + if "-" not in integrity: + return None + algo, b64 = integrity.split("-", 1) + algo = algo.strip().lower() + if algo not in ("sha256", "sha384", "sha512"): + return None + try: + digest = _b64.b64decode(b64, validate = True) + except Exception: + return None + return algo, digest + + +def download_tarball( + entry: PackageEntry, + dest: Path, + *, + timeout: float = HARD_HTTP_TIMEOUT_S, + max_bytes: int = HARD_MAX_TARBALL_BYTES, +) -> tuple[Path, str | None]: + """Stream-download entry.resolved to dest. Verify SRI integrity. + + Returns (downloaded_path, error_or_none). On any error the + returned path may not exist. Network access is restricted to + https://{ALLOWED_DOWNLOAD_HOST}/ -- the caller passes a Request + we already validated. + """ + # Re-assert hostname; the entry was validated at parse time but a + # defence-in-depth check here means a future refactor cannot + # accidentally bypass it. + parsed = urllib.parse.urlparse(entry.resolved) + if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST: + return dest, (f"refused download from non-allowlisted URL {entry.resolved!r}") + + decoded = _decode_integrity(entry.integrity or "") + if decoded is None: + return dest, f"unparseable integrity field {entry.integrity!r}" + algo, expected_digest = decoded + h = hashlib.new(algo) + + req = urllib.request.Request( + entry.resolved, + headers = { + "User-Agent": "unsloth-scan-npm-packages/1.0 (+supply-chain audit)", + "Accept": "application/octet-stream", + }, + method = "GET", + ) + try: + with urllib.request.urlopen(req, timeout = timeout) as r: + # Advertised length, if any. + cl = r.headers.get("Content-Length") + if cl is not None: + try: + cl_int = int(cl) + if cl_int > max_bytes: + return dest, (f"Content-Length {cl_int} > cap {max_bytes}") + except ValueError: + pass + written = 0 + with open(dest, "wb") as out: + while True: + chunk = r.read(64 * 1024) + if not chunk: + break + written += len(chunk) + if written > max_bytes: + return dest, ( + f"download exceeded cap {max_bytes} bytes " + f"after {written} bytes" + ) + h.update(chunk) + out.write(chunk) + except Exception as exc: + return dest, f"download failed: {exc}" + + actual = h.digest() + if actual != expected_digest: + return dest, ( + f"integrity mismatch: expected {algo}={_b64.b64encode(expected_digest).decode()!r}, " + f"got {algo}={_b64.b64encode(actual).decode()!r}" + ) + return dest, None + + +# ───────────────────────────────────────────────────────────────────── +# Safe tar extraction. Every Tarfile member is policed before write. +# ───────────────────────────────────────────────────────────────────── + + +def _is_within(root: Path, candidate: Path) -> bool: + try: + return candidate.resolve().is_relative_to(root.resolve()) + except (AttributeError, ValueError): + # Python <3.9 fallback (we target 3.10+ but be defensive). + try: + candidate.resolve().relative_to(root.resolve()) + return True + except Exception: + return False + + +def safe_extract( + tarball_path: Path, + extract_root: Path, + *, + max_total_bytes: int = HARD_MAX_TOTAL_BYTES, + max_members: int = HARD_MAX_MEMBERS, +) -> str | None: + """Extract tarball_path under extract_root with policed members. + + Returns None on success, or a string describing the refusal. + Streams via `r|gz` so we can abort mid-extraction without having + materialised the rest of the archive. + """ + extract_root.mkdir(parents = True, exist_ok = True) + total = 0 + count = 0 + try: + # Open in streaming mode so we never seek backwards in the + # input. `r|gz` rejects malformed gzip frames immediately. + with tarfile.open(tarball_path, mode = "r|gz") as tf: + for member in tf: + count += 1 + if count > max_members: + return f"member count {count} exceeded cap {max_members}" + name = member.name + # Reject obvious path-escape. + if name.startswith("/") or ".." in Path(name).parts: + return f"refused unsafe member name {name!r}" + # Reject device files, FIFOs, sockets, symlinks, hardlinks. + if member.issym() or member.islnk(): + return f"refused link member {name!r} (sym/lnk)" + if member.isdev() or member.isfifo(): + return f"refused special member {name!r}" + # Cumulative cap is checked against DECLARED size up + # front to short-circuit obvious bombs without reading + # the body. + declared = max(member.size, 0) + if declared > HARD_MAX_BINARY_FILE_BYTES: + return ( + f"member {name!r} declared size {declared} > " + f"absolute cap {HARD_MAX_BINARY_FILE_BYTES}" + ) + if total + declared > max_total_bytes: + return ( + f"cumulative bytes {total + declared} > cap " + f"{max_total_bytes} at {name!r}" + ) + # Strip leading "package/" -- the npm convention. We do + # NOT trust npm to be right, so we explicitly resolve + # the destination and refuse anything that escapes. + dest = extract_root / name + if not _is_within(extract_root, dest): + return f"refused escape: {name!r} resolved outside root" + if member.isdir(): + dest.mkdir(parents = True, exist_ok = True) + continue + if not member.isfile(): + # Anything we didn't classify above is unknown. + return f"refused unknown member type for {name!r}" + dest.parent.mkdir(parents = True, exist_ok = True) + src = tf.extractfile(member) + if src is None: + continue + # Sniff first 16 bytes to classify text vs binary. + # Text-cap members get the tight 16 MiB limit; binary + # members (executables, .node, .wasm, native libs) + # get the generous binary cap. We bound BOTH cases. + header = src.read(16) + is_binary = _looks_binary(name, header) + file_cap = ( + HARD_MAX_BINARY_FILE_BYTES + if is_binary + else HARD_MAX_TEXT_FILE_BYTES + ) + if declared > file_cap: + return ( + f"member {name!r} declared size {declared} > " + f"cap {file_cap} ({'binary' if is_binary else 'text'})" + ) + # Read remainder, bounded. + remainder_cap = file_cap - len(header) + rest = src.read(remainder_cap + 1) + data = header + rest + if len(data) > file_cap: + return ( + f"member {name!r} body exceeded declared size cap " + f"({'binary' if is_binary else 'text'})" + ) + total += len(data) + # Write with restrictive mode (rw-r--r--) so even if + # someone runs the extract dir nothing is executable. + with open(dest, "wb") as out: + out.write(data) + os.chmod(dest, 0o644) + except tarfile.TarError as exc: + return f"tar parse error: {exc}" + except Exception as exc: + return f"unexpected extract error: {exc!r}" + return None + + +# ───────────────────────────────────────────────────────────────────── +# Content scanning. +# ───────────────────────────────────────────────────────────────────── + + +def _evidence(text: str, pat: re.Pattern, max_chars: int = 200) -> str: + m = pat.search(text) + if not m: + return "" + start = max(0, m.start() - 30) + end = min(len(text), m.end() + 30) + snippet = text[start:end].replace("\n", " ") + if len(snippet) > max_chars: + snippet = snippet[:max_chars] + "..." + return snippet + + +LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare") + + +def scan_package_json( + pkg: PackageEntry, + rel: str, + text: str, +) -> list[Finding]: + findings: list[Finding] = [] + try: + meta = json.loads(text) + except Exception: + return findings + if not isinstance(meta, dict): + return findings + scripts = meta.get("scripts") or {} + if not isinstance(scripts, dict): + return findings + for hook in LIFECYCLE_HOOKS: + body = scripts.get(hook) + if not isinstance(body, str): + continue + if _LIFECYCLE_FETCH_EXEC.search(body): + findings.append( + Finding( + severity = CRITICAL, + package = pkg.display, + filename = rel, + pattern = f"lifecycle-fetch-exec ({hook})", + evidence = body, + detail = ( + f"`scripts.{hook}` fetches an external " + "resource and pipes/chains it to an " + "interpreter; this is the install-time RCE " + "vector. Refusing to install." + ), + ) + ) + # Credential file paths inside a lifecycle script are + # exfiltration prep -- npm runs these scripts automatically + # on `npm ci`. Manual `scripts.*` entries (like a `docker` + # dev script) are out of scope: npm does not run them. + for path_substr, why in CRED_PATH_SUBSTRINGS: + if path_substr in body: + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = f"cred-path-in-lifecycle ({hook})", + evidence = body, + detail = ( + f"`scripts.{hook}` references {why} " + f"({path_substr!r}); install-time access " + "to local credential files is the " + "exfiltration prep step" + ), + ) + ) + if _JS_ENV_TOKEN.search(body): + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = f"cred-env-in-lifecycle ({hook})", + evidence = _evidence(body, _JS_ENV_TOKEN), + detail = ( + f"`scripts.{hook}` references a credential " + "env var (GITHUB_TOKEN / NPM_TOKEN / AWS_* " + "/ etc); install-time access to runner " + "secrets is the exfiltration prep step" + ), + ) + ) + # Optional deps pointing at github: are the TanStack-style + # injection vector. + opt = meta.get("optionalDependencies") or {} + if isinstance(opt, dict): + for k, v in opt.items(): + if isinstance(v, str) and ( + v.startswith("github:") + or v.startswith("git+") + or v.startswith("git://") + ): + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = "optional-dep-non-registry", + evidence = f"{k}={v}", + detail = ( + "package.json `optionalDependencies` " + "points at a non-registry source; this " + "is the Shai-Hulud worm injection shape." + ), + ) + ) + return findings + + +def _host_in_outbound_context(text: str, host: str) -> bool: + """True if `host` appears in a way consistent with an outbound call. + + A bare `"169.254.169.254"` array literal (defensive blocklist) is + safe; a `fetch("http://169.254.169.254/...")` is not. The signal + is co-occurrence with either an HTTP URL scheme or a fetch verb + within a short window. + + A defensive blocklist looks like: + const CLOUD_METADATA_IPS = ["169.254.169.254", "169.254.170.2"]; + An exfil call looks like: + fetch("http://169.254.169.254/latest/meta-data/...") + http.request({ host: "169.254.169.254", path: "/..." }) + """ + # Esc for use in a regex (IPs contain dots). + host_re = re.escape(host) + # 1. URL form: http://host or https://host or //host/ or //host" + url_form = re.compile( + rf"(?:https?:)?//{host_re}(?:[:/\"'?#]|$)", + ) + if url_form.search(text): + return True + # 2. host appears within 200 chars of a fetch verb (either side). + fetch_context = re.compile( + rf"(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}{host_re}" + rf"|{host_re}[^\n]{{0,200}}(?:{_FETCH_VERBS_PAT})", + re.IGNORECASE, + ) + if fetch_context.search(text): + return True + # 3. `host:` / `hostname:` config field referencing the IP. + cfg_form = re.compile( + rf"(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`]", + re.IGNORECASE, + ) + if cfg_form.search(text): + return True + return False + + +def scan_text_blob( + pkg: PackageEntry, + rel: str, + text: str, +) -> list[Finding]: + findings: list[Finding] = [] + + # IOC substrings (literal, case-sensitive). + for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): + if needle in text: + findings.append( + Finding( + severity = sev, + package = pkg.display, + filename = rel, + pattern = "known-ioc-string", + evidence = needle, + detail = f"{why}: {needle!r}", + ) + ) + + # Credential surfaces. Tier 1: hosts with no legitimate use, + # bare substring is enough. + for needle, why in CRED_HOST_ALWAYS_BAD: + if needle in text: + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = "cred-surface-host (always-bad)", + evidence = needle, + detail = ( + f"references {why} ({needle!r}); no legitimate " + "frontend use of this surface" + ), + ) + ) + + # Credential surfaces. Tier 2: hosts that do appear in defensive + # code; require co-occurrence with a fetch verb or URL prefix. + for needle, why in CRED_HOST_NEEDS_CONTEXT: + if needle in text and _host_in_outbound_context(text, needle): + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = "cred-surface-host (outbound)", + evidence = needle, + detail = ( + f"references {why} ({needle!r}) in an outbound " + "call / URL / host config; a defensive blocklist " + "literal would not match this rule" + ), + ) + ) + + # Credential PATHS are deliberately not scanned here; they have + # too high a false-positive rate at file scope (defensive code, + # docker mounts, AWS SDK docs strings). `scan_package_json` + # catches the malicious case -- credential paths inside a + # lifecycle script run automatically on `npm ci`. + + # JS-specific regex. + if _JS_FETCH_EVAL.search(text): + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = "js-fetch-eval", + evidence = _evidence(text, _JS_FETCH_EVAL), + detail = ( + "Function/eval against base64-decoded payload " + "(obfuscated dropper shape)" + ), + ) + ) + if _JS_ENV_TOKEN.search(text): + findings.append( + Finding( + severity = MEDIUM, + package = pkg.display, + filename = rel, + pattern = "js-env-token", + evidence = _evidence(text, _JS_ENV_TOKEN), + detail = ("references credential env vars in package source"), + ) + ) + if _OBFUSC_BLOB.search(text): + findings.append( + Finding( + severity = HIGH, + package = pkg.display, + filename = rel, + pattern = "obfuscated-blob", + evidence = _evidence(text, _OBFUSC_BLOB), + detail = ( + "large base64-ish blob fed to Function/eval; " + "matches the TanStack worm dropper shape" + ), + ) + ) + + return findings + + +# Filename suffix decides which scanners run. We deliberately treat +# *.cjs/*.mjs/*.ts the same as *.js -- attackers use whichever +# extension the consumer's bundler / loader resolves. +_TEXT_SUFFIXES = ( + ".js", + ".mjs", + ".cjs", + ".ts", + ".tsx", + ".json", + ".html", + ".htm", + ".sh", + ".bash", + ".zsh", + ".py", + ".rb", + ".yml", + ".yaml", +) + + +def scan_extracted_tree( + pkg: PackageEntry, + root: Path, +) -> list[Finding]: + findings: list[Finding] = [] + for path in sorted(root.rglob("*")): + if not path.is_file(): + continue + rel = path.relative_to(root).as_posix() + lower = rel.lower() + if not lower.endswith(_TEXT_SUFFIXES): + # Skip native binaries entirely -- regex over compiled + # machine code is just noise (false positives in WASM + # opcodes, .node BSS segments, image pixel data). Use + # content-magic detection so extensionless executables + # (eg `package/biome`) and versioned shared libraries + # are also skipped. + try: + if path.stat().st_size > HARD_MAX_TEXT_FILE_BYTES: + continue + with open(path, "rb") as fh: + header = fh.read(16) + if _looks_binary(rel, header): + continue + data = header + path.read_bytes()[len(header) :] + except OSError: + continue + text = data.decode("utf-8", errors = "replace") + for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): + if needle in text: + findings.append( + Finding( + severity = sev, + package = pkg.display, + filename = rel, + pattern = "known-ioc-string", + evidence = needle, + detail = f"{why}: {needle!r}", + ) + ) + continue + try: + data = path.read_bytes() + except OSError: + continue + text = data.decode("utf-8", errors = "replace") + if rel.endswith("package.json"): + findings.extend(scan_package_json(pkg, rel, text)) + findings.extend(scan_text_blob(pkg, rel, text)) + return findings + + +# ───────────────────────────────────────────────────────────────────── +# Orchestrator. +# ───────────────────────────────────────────────────────────────────── + + +def scan_one( + pkg: PackageEntry, + workspace: Path, +) -> tuple[list[Finding], str | None]: + """Download + extract + scan a single package. Cleans up its dir. + + Returns (findings, error). `error` is non-None only on hard + failures (download error, integrity mismatch, malformed tarball); + on a clean run with findings the error is None and the caller + decides exit code based on severity. + """ + pkg_dir = workspace / f"{pkg.name.replace('/', '_')}-{pkg.version}" + pkg_dir.mkdir(parents = True, exist_ok = True) + tarball = pkg_dir / "pkg.tgz" + extract = pkg_dir / "x" + try: + _, err = download_tarball(pkg, tarball) + if err: + return [], err + err = safe_extract(tarball, extract) + if err: + return [], err + return scan_extracted_tree(pkg, extract), None + finally: + # Always wipe per-package data to keep the workspace bounded. + try: + shutil.rmtree(pkg_dir, ignore_errors = True) + except Exception: + pass + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description = "Pre-install npm tarball content scanner.", + ) + parser.add_argument( + "--lockfile", + default = str(REPO_ROOT / "studio" / "frontend" / "package-lock.json"), + help = "Path to package-lock.json (default: studio/frontend).", + ) + parser.add_argument( + "--max-packages", + type = int, + default = 0, + help = ( + "Cap on number of packages to scan (0 = no cap). Useful " + "for local triage; CI runs with 0." + ), + ) + parser.add_argument( + "--fail-on", + choices = ("info", "medium", "high", "critical"), + default = "high", + help = ( + "Lowest severity that fails the run (default: high). " + "Medium and below print but exit 0." + ), + ) + args = parser.parse_args(argv) + + lockfile = Path(args.lockfile).resolve() + if not lockfile.exists(): + print(f"[scan-npm] lockfile not found: {lockfile}", file = sys.stderr) + return 2 + + entries, struct_findings = parse_lockfile(lockfile) + if struct_findings: + print( + f"[scan-npm] {len(struct_findings)} structural finding(s) " + "from lockfile pass; subsequent download scan skipped for " + "those entries.", + flush = True, + ) + + if args.max_packages > 0: + entries = entries[: args.max_packages] + + workspace = Path(tempfile.mkdtemp(prefix = "npm-scan-")).resolve() + atexit.register(lambda: shutil.rmtree(workspace, ignore_errors = True)) + print( + f"[scan-npm] workspace: {workspace}\n" + f"[scan-npm] scanning {len(entries)} package(s) from {lockfile}", + flush = True, + ) + + all_findings: list[Finding] = list(struct_findings) + hard_errors: list[tuple[str, str]] = [] + + for i, pkg in enumerate(entries, start = 1): + print( + f"[scan-npm] [{i}/{len(entries)}] {pkg.display}", + flush = True, + ) + findings, err = scan_one(pkg, workspace) + if err: + hard_errors.append((pkg.display, err)) + print(f"[scan-npm] ERROR {pkg.display}: {err}", flush = True) + continue + all_findings.extend(findings) + for f in findings: + print(str(f), flush = True) + + # Sort by severity then package. + all_findings.sort(key = lambda f: (_SEVERITY_RANK[f.severity], f.package)) + + print( + f"\n[scan-npm] summary: {len(entries)} package(s), " + f"{len(all_findings)} finding(s), " + f"{len(hard_errors)} hard error(s)", + flush = True, + ) + + if hard_errors: + print("\n[scan-npm] HARD ERRORS:", file = sys.stderr) + for pkg, err in hard_errors: + print(f" {pkg}: {err}", file = sys.stderr) + + threshold = { + "info": INFO, + "medium": MEDIUM, + "high": HIGH, + "critical": CRITICAL, + }[args.fail_on] + threshold_rank = _SEVERITY_RANK[threshold] + blocking = [f for f in all_findings if _SEVERITY_RANK[f.severity] <= threshold_rank] + if hard_errors or blocking: + if blocking: + print( + f"\n[scan-npm] FAIL: {len(blocking)} finding(s) " + f"at or above {threshold}", + file = sys.stderr, + ) + return 1 + print("\n[scan-npm] OK", flush = True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/stamp_studio_release.py b/scripts/stamp_studio_release.py new file mode 100755 index 0000000000..5547b95137 --- /dev/null +++ b/scripts/stamp_studio_release.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Stamp and verify display-only Studio release metadata for builds.""" + +from __future__ import annotations + +import argparse +import os +import re +import subprocess +import sys +import tarfile +import zipfile +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +BUILD_INFO_PATH = ( + REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py" +) +BUILD_INFO_SUFFIX = "studio/backend/utils/_studio_release_build.py" +VERSION_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$") +GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$") +MAX_VERSION_LENGTH = 64 +PLACEHOLDER = """# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +\"\"\"Build-stamped Studio release metadata. + +Release builds may rewrite this module in the build workspace before creating +Python artifacts. Keep the committed value neutral so source checkouts do not +accidentally report a stale release tag. +\"\"\" + +STUDIO_RELEASE_VERSION = None +""" + + +def is_valid_version(value: object) -> bool: + if not isinstance(value, str): + return False + version = value.strip() + if not version or len(version) > MAX_VERSION_LENGTH: + return False + if version.endswith("-dirty") or GIT_DESCRIBE_SUFFIX_RE.search(version): + return False + return VERSION_RE.fullmatch(version) is not None + + +def _exact_git_tag() -> str | None: + try: + result = subprocess.run( + [ + "git", + "describe", + "--tags", + "--exact-match", + "--match", + "v[0-9]*", + "HEAD", + ], + cwd = REPO_ROOT, + check = False, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 2.0, + ) + except (OSError, subprocess.TimeoutExpired): + return None + if result.returncode != 0: + return None + tag = result.stdout.strip() + return tag if is_valid_version(tag) else None + + +def _git_worktree_is_dirty() -> bool: + try: + result = subprocess.run( + ["git", "status", "--porcelain"], + cwd = REPO_ROOT, + check = False, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 2.0, + ) + except (OSError, subprocess.TimeoutExpired): + return True + if result.returncode != 0: + return True + return bool(result.stdout.strip()) + + +def _github_tag() -> str | None: + if os.environ.get("GITHUB_REF_TYPE") != "tag": + return None + github_ref = os.environ.get("GITHUB_REF_NAME", "").strip() + return github_ref or None + + +def resolve_version() -> tuple[str | None, str]: + env_version = os.environ.get("UNSLOTH_STUDIO_RELEASE_VERSION", "").strip() + if env_version: + return (env_version, "UNSLOTH_STUDIO_RELEASE_VERSION") + + github_ref = _github_tag() + if github_ref: + return (github_ref, "GITHUB_REF_NAME") + + git_tag = _exact_git_tag() + if git_tag: + return (git_tag, "exact git tag") + + return (None, "none") + + +def build_info_source(version: str | None) -> str: + literal = repr(version) if version is not None else "None" + return f'''# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Build-stamped Studio release metadata.""" + +STUDIO_RELEASE_VERSION = {literal} +''' + + +def _env_version_conflicts(version: str) -> list[tuple[str, str]]: + conflicts: list[tuple[str, str]] = [] + github_ref = _github_tag() + if github_ref and is_valid_version(github_ref) and github_ref != version: + conflicts.append(("GITHUB_REF_NAME", github_ref)) + + git_tag = _exact_git_tag() + if git_tag and git_tag != version: + conflicts.append(("exact git tag", git_tag)) + + return conflicts + + +def stamp(require_release: bool) -> int: + version, source = resolve_version() + if version is not None and not is_valid_version(version): + print( + f"Invalid Studio release version from {source}: {version!r}", + file = sys.stderr, + ) + return 2 + + if version is not None and source == "UNSLOTH_STUDIO_RELEASE_VERSION": + conflicts = _env_version_conflicts(version) + if conflicts: + details = ", ".join(f"{name}={value!r}" for name, value in conflicts) + print( + "UNSLOTH_STUDIO_RELEASE_VERSION does not match available " + f"release tag metadata: {details}", + file = sys.stderr, + ) + return 2 + + if require_release and source == "exact git tag" and _git_worktree_is_dirty(): + print( + "Refusing to publish from a dirty exact-tag checkout. Set " + "UNSLOTH_STUDIO_RELEASE_VERSION explicitly from release automation " + "or publish from a clean tag checkout.", + file = sys.stderr, + ) + return 2 + + if version is None: + if require_release: + print( + "No Studio release version available. Set " + "UNSLOTH_STUDIO_RELEASE_VERSION, build from a GitHub tag, " + "or run from an exact local Studio release tag.", + file = sys.stderr, + ) + return 2 + BUILD_INFO_PATH.write_text(PLACEHOLDER, encoding = "utf-8") + print("dev") + return 0 + + BUILD_INFO_PATH.write_text(build_info_source(version), encoding = "utf-8") + print(f"Stamping Studio release version {version} from {source}", file = sys.stderr) + print(version) + return 0 + + +def _read_wheel_member(path: Path) -> str | None: + with zipfile.ZipFile(path) as archive: + for name in archive.namelist(): + if name.endswith(BUILD_INFO_SUFFIX): + return archive.read(name).decode("utf-8") + return None + + +def _read_sdist_member(path: Path) -> str | None: + with tarfile.open(path) as archive: + for member in archive.getmembers(): + if member.name.endswith(BUILD_INFO_SUFFIX): + extracted = archive.extractfile(member) + if extracted is None: + return None + return extracted.read().decode("utf-8") + return None + + +def verify_dist(expected: str, dist_dir: Path) -> int: + if not is_valid_version(expected): + print(f"Invalid expected Studio release version: {expected!r}", file = sys.stderr) + return 2 + + artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz")) + if not artifacts: + print(f"No wheel or sdist artifacts found in {dist_dir}", file = sys.stderr) + return 2 + + expected_line = f"STUDIO_RELEASE_VERSION = {expected!r}" + failures: list[str] = [] + for artifact in artifacts: + if artifact.suffix == ".whl": + content = _read_wheel_member(artifact) + else: + content = _read_sdist_member(artifact) + if content is None: + failures.append(f"{artifact.name}: missing {BUILD_INFO_SUFFIX}") + elif expected_line not in content: + failures.append(f"{artifact.name}: Studio release version mismatch") + + if failures: + for failure in failures: + print(failure, file = sys.stderr) + return 2 + + print(f"Verified Studio release version {expected} in {len(artifacts)} artifact(s)") + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description = __doc__) + parser.add_argument("--require-release", action = "store_true") + parser.add_argument("--verify-dist", type = Path) + parser.add_argument("--expected") + args = parser.parse_args() + + if args.verify_dist is not None: + if not args.expected: + parser.error("--verify-dist requires --expected") + return verify_dist(args.expected, args.verify_dist) + + return stamp(require_release = args.require_release) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 38c0261f5a..6c97ef8cb2 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -956,6 +956,66 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: + """Return DLL dirs from pip-installed CUDA wheels under + ``/Lib/site-packages/`` so llama-server.exe can load + ``cudart64_X.dll`` / ``cublas64_X.dll`` without a system CUDA + toolkit. Mirrors the Linux ``nvidia/cu*/lib`` LD_LIBRARY_PATH + block, with parity for the Windows-specific wheel layouts seen + in the wild. Covered patterns: + * ``nvidia//bin`` -- legacy modular wheels + (``nvidia-cuda-runtime-cu12``, ``nvidia-cublas-cu12``, etc.). + * ``nvidia//bin/x86_64`` and ``.../bin/x64`` -- current + CUDA 13 wheel layout used by the unsuffixed + ``nvidia-cuda-runtime`` / ``nvidia-cublas`` packages, which + ship under ``nvidia/cu13/bin/x86_64/`` (#5106). + * ``nvidia//Library/bin`` (and arch subdirs) -- conda- + style wheel repacks. + * ``torch/lib`` -- PyTorch's own CUDA-bundled Windows wheel, + which can ship ``cudart64_*.dll`` directly here instead of + as separate ``nvidia-*`` wheels. The install-side helper + ``python_runtime_dirs`` in ``install_llama_prebuilt.py`` + covers this path for the same reason. + + Walks the tree with ``Path.iterdir`` rather than ``glob.glob`` + so the resolver is safe against Windows paths containing + ``[`` or ``]`` (valid in usernames; would otherwise be + interpreted as a glob character class and silently miss + existing dirs).""" + site_packages = Path(prefix) / "Lib" / "site-packages" + out: list[str] = [] + seen: set[str] = set() + + def _add(path: Path) -> None: + if not path.is_dir(): + return + key = os.path.normcase(os.path.abspath(str(path))) + if key in seen: + return + seen.add(key) + out.append(str(path)) + + nvidia_root = site_packages / "nvidia" + if nvidia_root.is_dir(): + for pkg_dir in nvidia_root.iterdir(): + if not pkg_dir.is_dir(): + continue + # Order matters for PATH search: arch-specific subdirs + # first so the explicit cudart64_X.dll location wins + # over a sibling ``bin`` that might be empty. + for sub in ( + pkg_dir / "bin" / "x86_64", + pkg_dir / "bin" / "x64", + pkg_dir / "bin", + pkg_dir / "Library" / "bin" / "x86_64", + pkg_dir / "Library" / "bin" / "x64", + pkg_dir / "Library" / "bin", + ): + _add(sub) + _add(site_packages / "torch" / "lib") + return out + @staticmethod def _select_gpus( model_size_bytes: int, @@ -2319,9 +2379,14 @@ class LlamaCppBackend: binary_dir = str(Path(binary).parent) if sys.platform == "win32": - # On Windows, CUDA DLLs (cublas64_12.dll, cudart64_12.dll, etc.) - # must be on PATH. Add CUDA_PATH\bin if available. + # CUDA DLLs (cudart64_X.dll, cublas64_X.dll, etc.) must + # be on PATH. Order: binary_dir, torch's pip-installed + # nvidia wheels, then a system CUDA toolkit. Pip wheels + # are the canonical source per Studio's install design + # (mirrors the Linux LD_LIBRARY_PATH block below) and + # CUDA_PATH covers users with a system toolkit. #5106. path_dirs = [binary_dir] + path_dirs.extend(self._windows_pip_nvidia_dll_dirs(sys.prefix)) cuda_path = os.environ.get("CUDA_PATH", "") if cuda_path: cuda_bin = os.path.join(cuda_path, "bin") diff --git a/studio/backend/main.py b/studio/backend/main.py index 420b0db435..43a55e36f2 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -135,6 +135,11 @@ import utils.hardware.hardware as _hw_module from utils.cache_cleanup import clear_unsloth_compiled_cache from utils.native_path_leases import native_path_leases_supported +from utils.update_status import ( + get_studio_install_source_status, + get_studio_update_status, +) +from utils.studio_version import get_studio_version def get_unsloth_version() -> str: @@ -156,6 +161,7 @@ def get_unsloth_version() -> str: UNSLOTH_VERSION = get_unsloth_version() +STUDIO_VERSION = get_studio_version() @asynccontextmanager @@ -303,6 +309,7 @@ async def health_check(): "timestamp": datetime.now().isoformat(), "service": "Unsloth UI Backend", "version": UNSLOTH_VERSION, + "studio_version": STUDIO_VERSION, "device_type": device_type, "chat_only": _hw_module.CHAT_ONLY, "desktop_protocol_version": 1, @@ -315,6 +322,18 @@ async def health_check(): } +@app.get("/api/studio/install-source") +def studio_install_source(_current_subject: str = Depends(get_current_subject)): + """Return source-aware install metadata without remote update checks.""" + return get_studio_install_source_status(UNSLOTH_VERSION) + + +@app.get("/api/studio/update-status") +def studio_update_status(_current_subject: str = Depends(get_current_subject)): + """Return source-aware manual update status for browser-served Studio.""" + return get_studio_update_status(UNSLOTH_VERSION) + + @app.post("/api/shutdown") async def shutdown_server( request: Request, diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 569f6eff3d..96f8816b57 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -3,6 +3,7 @@ typer fastapi uvicorn pydantic +packaging matplotlib pandas nest_asyncio diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py new file mode 100644 index 0000000000..7d4719c0e7 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the Windows pip-nvidia DLL dir resolver. + +Studio installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13, +nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find +those DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH +block. See unslothai/unsloth#5106. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub heavy deps before importing the module under test. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + + +def _make_nvidia_layout(prefix: Path, pkgs_with_layout: dict[str, str]): + """Build a fake /Lib/site-packages/nvidia//{bin|Library/bin} + tree with a stub DLL inside each leaf so isdir() picks them up.""" + nv = prefix / "Lib" / "site-packages" / "nvidia" + for pkg, layout in pkgs_with_layout.items(): + if layout == "bin": + d = nv / pkg / "bin" + elif layout == "library_bin": + d = nv / pkg / "Library" / "bin" + else: + raise ValueError(layout) + d.mkdir(parents = True, exist_ok = True) + (d / "stub.dll").write_bytes(b"") + + +class TestWindowsPipNvidiaDllDirs: + def test_returns_empty_when_no_nvidia_wheels(self, tmp_path): + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_picks_up_bin_layout(self, tmp_path): + _make_nvidia_layout( + tmp_path, + { + "cuda_runtime": "bin", + "cublas": "bin", + "cudnn": "bin", + }, + ) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 3 + assert all(Path(p).is_dir() for p in result) + assert all(Path(p).name == "bin" for p in result) + names = {Path(p).parent.name for p in result} + assert names == {"cuda_runtime", "cublas", "cudnn"} + + def test_picks_up_library_bin_layout(self, tmp_path): + _make_nvidia_layout( + tmp_path, + { + "cuda_runtime": "library_bin", + "cublas": "library_bin", + }, + ) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 2 + for p in result: + assert Path(p).is_dir() + assert Path(p).parent.name == "Library" + assert Path(p).parent.parent.name in {"cuda_runtime", "cublas"} + + def test_mixed_layouts_all_resolved(self, tmp_path): + _make_nvidia_layout( + tmp_path, + { + "cuda_runtime": "bin", + "cublas": "library_bin", + "cudnn": "bin", + "nvjitlink": "library_bin", + }, + ) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 4 + + def test_does_not_walk_outside_known_paths(self, tmp_path): + # Only nvidia//{bin,Library/bin} and torch/lib are picked + # up. Unrelated site-packages contents (numpy, scipy, ...) must + # be ignored. + site = tmp_path / "Lib" / "site-packages" + (site / "numpy").mkdir(parents = True) + (site / "scipy" / "linalg").mkdir(parents = True) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_picks_up_torch_lib(self, tmp_path): + # PyTorch's Windows CUDA wheel bundles cudart64_X.dll / + # cublas64_X.dll directly under Lib/site-packages/torch/lib/ + # instead of as separate nvidia-* wheels. Without this, users + # on torch-bundled-CUDA installs still hit #5106. + torch_lib = tmp_path / "Lib" / "site-packages" / "torch" / "lib" + torch_lib.mkdir(parents = True) + (torch_lib / "cudart64_12.dll").write_bytes(b"") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 1 + assert Path(result[0]) == torch_lib + + def test_torch_lib_combined_with_nvidia_wheels(self, tmp_path): + # Both modular nvidia-* wheels and torch/lib are returned when + # present together. + _make_nvidia_layout( + tmp_path, + { + "cuda_runtime": "bin", + "cublas": "bin", + }, + ) + torch_lib = tmp_path / "Lib" / "site-packages" / "torch" / "lib" + torch_lib.mkdir(parents = True) + (torch_lib / "cudart64_13.dll").write_bytes(b"") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert len(result) == 3 + names = {Path(p).name for p in result} + assert names == {"bin", "lib"} + assert any(Path(p) == torch_lib for p in result) + + def test_torch_lib_must_be_a_directory(self, tmp_path): + # If torch/lib exists as a file (broken install), it is + # ignored, not returned. + site = tmp_path / "Lib" / "site-packages" / "torch" + site.mkdir(parents = True) + (site / "lib").write_bytes(b"not a dir") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_skips_non_directories(self, tmp_path): + nv = tmp_path / "Lib" / "site-packages" / "nvidia" + (nv / "cuda_runtime").mkdir(parents = True) + # Create a regular file at the path where 'bin' would normally be a dir + (nv / "cuda_runtime" / "bin").write_bytes(b"not a dir") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert result == [] + + def test_missing_prefix_does_not_raise(self): + # If sys.prefix points to a path that doesn't exist (unusual, + # but possible during test setup), the resolver must just + # return [] rather than raising. + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs( + "/this/path/does/not/exist/anywhere" + ) + assert result == [] + + def test_picks_up_cu13_bin_x86_64_layout(self, tmp_path): + # Current ``nvidia-cuda-runtime`` 13.x and ``nvidia-cublas`` + # 13.x Windows wheels ship DLLs under + # ``nvidia/cu13/bin/x86_64/`` instead of ``nvidia//bin/``. + # Without this, users on the new CUDA 13 wheel generation hit + # the original #5106 failure mode. + dll_dir = ( + tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64" + ) + dll_dir.mkdir(parents = True) + for name in ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"): + (dll_dir / name).write_bytes(b"") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert str(dll_dir) in result, f"cu13 bin/x86_64 not in {result}" + + def test_picks_up_bin_x64_layout(self, tmp_path): + # Some repackaged wheels use ``bin/x64`` (Windows-x64 convention) + # instead of ``bin/x86_64`` (NVIDIA-internal convention). + dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x64" + dll_dir.mkdir(parents = True) + (dll_dir / "cudart64_13.dll").write_bytes(b"") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + assert str(dll_dir) in result + + def test_mixed_cu12_and_cu13_layouts(self, tmp_path): + # A venv could have both the modular cu12 wheels (legacy) and + # the unsuffixed cu13 wheel installed side by side. Both must + # be reachable. + site = tmp_path / "Lib" / "site-packages" + cu12_bin = site / "nvidia" / "cuda_runtime" / "bin" + cu13_arch = site / "nvidia" / "cu13" / "bin" / "x86_64" + cu12_bin.mkdir(parents = True) + cu13_arch.mkdir(parents = True) + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + result_set = {Path(p) for p in result} + assert cu12_bin in result_set + assert cu13_arch in result_set + + def test_glob_meta_in_prefix_is_safe(self, tmp_path): + # Windows usernames / install paths can contain ``[`` or ``]``. + # A glob-based resolver would interpret these as a character + # class and silently return [] even when DLL dirs exist. The + # iterdir-based implementation must work on such paths. + prefix = tmp_path / "studio_[gpu]_install" + dll_dir = prefix / "Lib" / "site-packages" / "nvidia" / "cuda_runtime" / "bin" + dll_dir.mkdir(parents = True) + (dll_dir / "cudart64_12.dll").write_bytes(b"") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix)) + assert str(dll_dir) in result, f"bracket-prefixed path returned empty: {result}" + + def test_arch_subdir_listed_before_parent_bin(self, tmp_path): + # When both ``nvidia//bin/`` and + # ``nvidia//bin/x86_64/`` exist, the arch-specific subdir + # must be listed first so Windows DLL search picks up the + # cudart64_X.dll location even if the parent ``bin`` is empty. + site = tmp_path / "Lib" / "site-packages" + outer_bin = site / "nvidia" / "cu13" / "bin" + arch_bin = outer_bin / "x86_64" + arch_bin.mkdir(parents = True) + (arch_bin / "cudart64_13.dll").write_bytes(b"") + result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path)) + # outer_bin exists as a directory (it contains arch_bin); the + # arch-specific subdir should come first in the list. + result_paths = [Path(p) for p in result] + assert arch_bin in result_paths + assert outer_bin in result_paths + assert result_paths.index(arch_bin) < result_paths.index(outer_bin) diff --git a/studio/backend/utils/_studio_release_build.py b/studio/backend/utils/_studio_release_build.py new file mode 100644 index 0000000000..267197a202 --- /dev/null +++ b/studio/backend/utils/_studio_release_build.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Build-stamped Studio release metadata. + +Release builds may rewrite this module in the build workspace before creating +Python artifacts. Keep the committed value neutral so source checkouts do not +accidentally report a stale release tag. +""" + +STUDIO_RELEASE_VERSION = None diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py new file mode 100644 index 0000000000..70059f8a3c --- /dev/null +++ b/studio/backend/utils/studio_version.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Network-free Studio release version resolution for display-only UI.""" + +from __future__ import annotations + +import re +import subprocess +from pathlib import Path + +from utils import _studio_release_build + +_DEV_VERSION = "dev" +_GIT_TIMEOUT_SECONDS = 1.0 +_STUDIO_TAG_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$") +_GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$") +_MAX_VERSION_LENGTH = 64 + + +def is_valid_studio_release_version(value: object) -> bool: + """Return True for Studio release tags such as ``v0.1.39-beta``.""" + if not isinstance(value, str): + return False + version = value.strip() + if not version or len(version) > _MAX_VERSION_LENGTH: + return False + if version.endswith("-dirty") or _GIT_DESCRIBE_SUFFIX_RE.search(version): + return False + return _STUDIO_TAG_RE.fullmatch(version) is not None + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[3] + + +def _path_is_in_site_packages(path: Path) -> bool: + return any(part in {"site-packages", "dist-packages"} for part in path.parts) + + +def _is_source_checkout(repo_root: Path) -> bool: + return (repo_root / ".git").exists() and not _path_is_in_site_packages( + Path(__file__).resolve() + ) + + +def _exact_git_studio_tag(repo_root: Path) -> str | None: + try: + result = subprocess.run( + [ + "git", + "describe", + "--tags", + "--exact-match", + "--match", + "v[0-9]*", + "HEAD", + ], + cwd = repo_root, + check = False, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = _GIT_TIMEOUT_SECONDS, + ) + except (OSError, subprocess.TimeoutExpired): + return None + + if result.returncode != 0: + return None + + tag = result.stdout.strip() + return tag if is_valid_studio_release_version(tag) else None + + +def get_studio_version(repo_root: Path | None = None) -> str: + """Return the installed Studio release tag for display, or ``dev``. + + This value is intentionally separate from the PyPI ``unsloth`` package + version used by update checks. It never performs network requests. + """ + resolved_repo_root = repo_root or _repo_root() + + if _is_source_checkout(resolved_repo_root): + git_tag = _exact_git_studio_tag(resolved_repo_root) + return git_tag if git_tag is not None else _DEV_VERSION + + stamped_version = _studio_release_build.STUDIO_RELEASE_VERSION + if is_valid_studio_release_version(stamped_version): + return stamped_version.strip() + + return _DEV_VERSION diff --git a/studio/backend/utils/update_status.py b/studio/backend/utils/update_status.py new file mode 100644 index 0000000000..9142203a69 --- /dev/null +++ b/studio/backend/utils/update_status.py @@ -0,0 +1,374 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Web update status helpers for browser-served Unsloth Studio. + +This module is intentionally side-effect light: no network work happens at +import time or from /api/health. The PyPI check is lazy, cached, and only used +for normal PyPI-managed installs. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +import urllib.request +from dataclasses import dataclass +from datetime import datetime, timezone +from importlib.metadata import PackageNotFoundError, distribution +from pathlib import Path +from typing import Any + +from packaging.version import InvalidVersion, Version + +PACKAGE_NAME = "unsloth" +PYPI_JSON_URL = "https://pypi.org/pypi/unsloth/json" +PYPI_TIMEOUT_SECONDS = 3 +PYPI_RESPONSE_MAX_BYTES = 5 * 1024 * 1024 +PYPI_SUCCESS_TTL_SECONDS = 12 * 60 * 60 +PYPI_FAILURE_TTL_SECONDS = 60 * 60 +RELEASE_NOTES_URL = "https://unsloth.ai/docs/new/changelog" +DISABLE_ENV_VAR = "UNSLOTH_DISABLE_UPDATE_CHECK" + +LOCAL_INSTALL_SOURCES = {"editable", "local_path", "vcs", "local_repo"} + + +@dataclass(frozen = True) +class LatestVersionResult: + latest_version: str | None + checked_at: str + reason: str | None = None + error: str | None = None + + +@dataclass +class _LatestVersionCacheEntry: + result: LatestVersionResult + expires_at: float + + +_cache_condition = threading.Condition() +_latest_version_cache: _LatestVersionCacheEntry | None = None +_latest_version_fetching = False + + +def reset_update_status_cache() -> None: + """Clear the in-process PyPI cache. Intended for tests.""" + global _latest_version_cache, _latest_version_fetching + with _cache_condition: + _latest_version_cache = None + _latest_version_fetching = False + _cache_condition.notify_all() + + +def detect_install_source() -> str: + """Return a coarse install source without exposing local paths. + + Sources are intentionally conservative. PEP 610 local/vcs metadata wins. + Legacy source installs are treated as local only when package files resolve + outside site-packages/dist-packages and under a Git checkout. + """ + try: + dist = distribution(PACKAGE_NAME) + except PackageNotFoundError: + return ( + "local_repo" + if _path_has_git_parent(_repo_root_from_this_file()) + else "unknown" + ) + + try: + direct_url = dist.read_text("direct_url.json") + except Exception: + return "unknown" + if direct_url: + return _source_from_direct_url(direct_url) + + for package_path in _distribution_package_paths(dist): + if not _path_is_under_python_package_dir(package_path) and _path_has_git_parent( + package_path + ): + return "local_repo" + + return "pypi" + + +def get_studio_install_source_status(current_version: str) -> dict[str, Any]: + """Return install-source metadata without remote update checks.""" + install_source = detect_install_source() + reason = None + if install_source in LOCAL_INSTALL_SOURCES: + reason = "local_source" + elif install_source == "unknown": + reason = "unknown_source" + + return _status_response( + current_version = current_version, + latest_version = None, + install_source = install_source, + reason = reason, + ) + + +def get_studio_update_status(current_version: str) -> dict[str, Any]: + """Return public, read-only update status for the web UI.""" + install_source = detect_install_source() + + if os.environ.get(DISABLE_ENV_VAR) == "1": + return _status_response( + current_version = current_version, + latest_version = None, + install_source = install_source, + reason = "disabled", + ) + + if install_source in LOCAL_INSTALL_SOURCES: + return _status_response( + current_version = current_version, + latest_version = None, + install_source = install_source, + reason = "local_source", + ) + + if install_source != "pypi": + return _status_response( + current_version = current_version, + latest_version = None, + install_source = install_source, + reason = "unknown_source", + ) + + current = _parse_current_version(current_version) + if current is None: + return _status_response( + current_version = current_version, + latest_version = None, + install_source = install_source, + reason = "invalid_current_version" + if current_version != "dev" + else "dev_build", + ) + latest_result = get_latest_pypi_version() + if latest_result.latest_version is None: + return _status_response( + current_version = current_version, + latest_version = None, + install_source = install_source, + reason = latest_result.reason or "offline", + error = latest_result.error, + checked_at = latest_result.checked_at, + ) + + try: + latest = Version(latest_result.latest_version) + except InvalidVersion: + return _status_response( + current_version = current_version, + latest_version = latest_result.latest_version, + install_source = install_source, + reason = "invalid_latest_version", + error = "PyPI returned an invalid version.", + checked_at = latest_result.checked_at, + ) + + if latest > current: + return _status_response( + current_version = current_version, + latest_version = latest_result.latest_version, + install_source = install_source, + update_available = True, + can_show_web_notification = True, + checked_at = latest_result.checked_at, + ) + + return _status_response( + current_version = current_version, + latest_version = latest_result.latest_version, + install_source = install_source, + reason = "current_not_older", + checked_at = latest_result.checked_at, + ) + + +def get_latest_pypi_version() -> LatestVersionResult: + """Return the latest PyPI version using a small in-process TTL cache.""" + global _latest_version_cache, _latest_version_fetching + + while True: + now = time.monotonic() + with _cache_condition: + if _latest_version_cache and _latest_version_cache.expires_at > now: + return _latest_version_cache.result + if not _latest_version_fetching: + _latest_version_fetching = True + break + _cache_condition.wait(timeout = PYPI_TIMEOUT_SECONDS + 1) + + try: + result = _fetch_latest_pypi_version() + except Exception: + result = LatestVersionResult( + latest_version = None, + checked_at = _utc_now_iso(), + reason = "offline", + error = "Could not check PyPI update metadata.", + ) + + ttl = ( + PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS + ) + with _cache_condition: + _latest_version_cache = _LatestVersionCacheEntry( + result = result, + expires_at = time.monotonic() + ttl, + ) + _latest_version_fetching = False + _cache_condition.notify_all() + return result + + +def _fetch_latest_pypi_version() -> LatestVersionResult: + checked_at = _utc_now_iso() + request = urllib.request.Request( + PYPI_JSON_URL, + headers = {"User-Agent": "unsloth-studio-update-check"}, + ) + + try: + with urllib.request.urlopen(request, timeout = PYPI_TIMEOUT_SECONDS) as response: + body = response.read(PYPI_RESPONSE_MAX_BYTES + 1) + if len(body) > PYPI_RESPONSE_MAX_BYTES: + return LatestVersionResult( + latest_version = None, + checked_at = checked_at, + reason = "malformed_response", + error = "PyPI returned oversized update metadata.", + ) + payload = json.loads(body.decode("utf-8")) + except json.JSONDecodeError: + return LatestVersionResult( + latest_version = None, + checked_at = checked_at, + reason = "malformed_response", + error = "PyPI returned malformed update metadata.", + ) + except OSError: + return LatestVersionResult( + latest_version = None, + checked_at = checked_at, + reason = "offline", + error = "Could not reach PyPI for update metadata.", + ) + + latest = ( + payload.get("info", {}).get("version") if isinstance(payload, dict) else None + ) + if not isinstance(latest, str) or not latest.strip(): + return LatestVersionResult( + latest_version = None, + checked_at = checked_at, + reason = "malformed_response", + error = "PyPI update metadata did not include a version.", + ) + + return LatestVersionResult(latest_version = latest.strip(), checked_at = checked_at) + + +def _status_response( + *, + current_version: str, + latest_version: str | None, + install_source: str, + reason: str | None = None, + error: str | None = None, + update_available: bool = False, + can_show_web_notification: bool = False, + checked_at: str | None = None, +) -> dict[str, Any]: + return { + "current_version": current_version, + "latest_version": latest_version, + "update_available": update_available, + "install_source": install_source, + "can_show_web_notification": can_show_web_notification, + "release_notes_url": RELEASE_NOTES_URL, + "checked_at": checked_at or _utc_now_iso(), + "reason": reason, + "error": error, + } + + +def _source_from_direct_url(direct_url: str) -> str: + try: + payload = json.loads(direct_url) + except json.JSONDecodeError: + return "unknown" + + if not isinstance(payload, dict): + return "unknown" + + dir_info = payload.get("dir_info") + if isinstance(dir_info, dict) and dir_info.get("editable") is True: + return "editable" + + if isinstance(payload.get("vcs_info"), dict): + return "vcs" + + url = payload.get("url") + if isinstance(url, str) and url.startswith("file:"): + return "local_path" + + return "unknown" + + +def _distribution_package_paths(dist: Any) -> list[Path]: + paths: list[Path] = [] + files = getattr(dist, "files", None) or [] + for file in files: + text = str(file) + if not text.startswith(("unsloth/", "unsloth_cli/", "studio/")): + continue + try: + paths.append(Path(dist.locate_file(file)).resolve()) + except OSError: + continue + return paths + + +def _path_is_under_python_package_dir(path: Path) -> bool: + return any(part in {"site-packages", "dist-packages"} for part in path.parts) + + +def _path_has_git_parent(path: Path) -> bool: + for candidate in (path, *path.parents): + if (candidate / ".git").exists(): + return True + return False + + +def _repo_root_from_this_file() -> Path: + # update_status.py -> utils -> backend -> studio -> repo root + try: + return Path(__file__).resolve().parents[3] + except IndexError: + return Path(__file__).resolve().parent + + +def _parse_current_version(current_version: str) -> Version | None: + if current_version == "dev": + return None + try: + return Version(current_version) + except InvalidVersion: + return None + + +def _utc_now_iso() -> str: + return ( + datetime.now(timezone.utc) + .replace(microsecond = 0) + .isoformat() + .replace("+00:00", "Z") + ) diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 62e78b809a..fe3173aa41 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -9,6 +9,7 @@ import { shouldUseCustomWindowTitlebar, } from "@/components/tauri/window-titlebar"; import { Toaster } from "@/components/ui/sonner"; +import { WebUpdateBanner } from "@/components/web/update-banner"; import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth"; import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain"; import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend"; @@ -154,6 +155,13 @@ const HIDDEN_TITLEBAR_SIDEBAR_ROUTES = new Set([ "/signup", ]); +const WEB_UPDATE_HIDDEN_ROUTES = new Set([ + "/onboarding", + "/login", + "/change-password", + "/signup", +]); + function TauriWrapper({ children }: { children: ReactNode }) { const pathname = useRouterState({ select: (s) => s.location.pathname }); const { @@ -234,7 +242,14 @@ function TauriWrapper({ children }: { children: ReactNode }) { return () => { disposed = true; }; }, [status, desktopAuthRetry]); - if (!isTauri) return <>{children}; + if (!isTauri) { + return ( + <> + {children} + + + ); + } const showApp = status === "running" && desktopAuthReady; const startupStatus = status === "running" ? "starting" : status; diff --git a/studio/frontend/src/components/web/update-banner.tsx b/studio/frontend/src/components/web/update-banner.tsx new file mode 100644 index 0000000000..685a8aa3dc --- /dev/null +++ b/studio/frontend/src/components/web/update-banner.tsx @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { useWebUpdateCheck } from "@/hooks/use-web-update-check"; +import { isTauri } from "@/lib/api-base"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { AnimatePresence, motion } from "motion/react"; +import { type ReactElement, useEffect, useRef, useState } from "react"; + +const STUDIO_UPDATE_CMD = "unsloth studio update"; +const RELEASE_NOTES_URL = "https://unsloth.ai/docs/new/changelog"; +const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1]; + +interface WebUpdateBannerProps { + enabled?: boolean; +} + +export function WebUpdateBanner({ + enabled = true, +}: WebUpdateBannerProps): ReactElement | null { + const { status, dismiss } = useWebUpdateCheck({ enabled }); + const [copiedVersion, setCopiedVersion] = useState(null); + const dismissTimerRef = useRef | null>(null); + + useEffect(() => { + return () => { + if (dismissTimerRef.current) { + clearTimeout(dismissTimerRef.current); + } + }; + }, []); + + if (isTauri) { + return null; + } + + async function handleCopyCommand() { + if (!(await copyToClipboard(STUDIO_UPDATE_CMD))) { + return; + } + setCopiedVersion(status?.latestVersion ?? null); + if (dismissTimerRef.current) { + clearTimeout(dismissTimerRef.current); + } + dismissTimerRef.current = setTimeout(() => dismiss(), 900); + } + + return ( + + {status ? ( + +
+ + +
+ +
+

+ Package update available: {status.latestVersion} +

+

+ Installed package: {status.currentVersion}. To update Studio, + run this in your terminal, then restart Studio. +

+
+
+ +
+ + + +
+
+ + ) : null} + + ); +} diff --git a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx index 4d0b87981b..e4cdccd2d7 100644 --- a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx +++ b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx @@ -1,8 +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 -import { cn } from "@/lib/utils"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { cn } from "@/lib/utils"; import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion, useReducedMotion } from "motion/react"; @@ -14,11 +14,42 @@ const STUDIO_UPDATE_FALLBACK_UNIX_CMD = "curl -fsSL https://unsloth.ai/install.sh | sh"; const STUDIO_UPDATE_FALLBACK_WINDOWS_CMD = "irm https://unsloth.ai/install.ps1 | iex"; +const STUDIO_LOCAL_PULL_CMD = "git pull --ff-only"; +const STUDIO_LOCAL_UPDATE_CMD = "unsloth studio update --local"; +const STUDIO_LOCAL_FALLBACK_UNIX_CMD = "./install.sh --local"; +const STUDIO_LOCAL_FALLBACK_WINDOWS_CMD = ".\\install.ps1 --local"; export type UpdateShell = "windows" | "unix"; +export type UpdateInstallSource = + | "pypi" + | "editable" + | "local_path" + | "vcs" + | "local_repo" + | "unknown"; +type UpdateInstallSourceState = UpdateInstallSource | "loading"; function getStudioUpdateInstructionLine(shell: UpdateShell): string { - return shell === "windows" ? "Open PowerShell and run:" : "Open Terminal and run:"; + return shell === "windows" + ? "Open PowerShell and run:" + : "Open Terminal and run:"; +} + +function isLocalInstallSource( + installSource?: UpdateInstallSourceState | null, +): boolean { + return Boolean( + installSource && + installSource !== "pypi" && + installSource !== "unknown" && + installSource !== "loading", + ); +} + +function isUnknownInstallSource( + installSource?: UpdateInstallSourceState | null, +): boolean { + return installSource === "unknown"; } function CopyableCommand({ @@ -54,7 +85,7 @@ function CopyableCommand({
{copied ? ( - + ) : ( )} @@ -77,24 +111,38 @@ function CopyableCommand({ ); } +// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: keep source-specific update guidance in one component so the command matrix stays visible. export function UpdateStudioInstructions({ className, defaultShell, + installSource, showTitle = true, }: { className?: string; defaultShell: UpdateShell; + installSource?: UpdateInstallSourceState | null; showTitle?: boolean; }): ReactElement { const [shell, setShell] = useState(defaultShell); const prefersReducedMotion = useReducedMotion(); const windows = shell === "windows"; + const localInstallSource = isLocalInstallSource(installSource); + const checkoutInstallSource = + installSource === "editable" || installSource === "local_repo"; + const packagedSourceInstall = + installSource === "vcs" || installSource === "local_path"; + const loadingInstallSource = installSource === "loading"; + const unknownInstallSource = isUnknownInstallSource(installSource); const fadeTransition = prefersReducedMotion ? { duration: 0 } : { duration: 0.16, ease: [0.165, 0.84, 0.44, 1] as const }; - const fadeInitial = prefersReducedMotion ? { opacity: 1 } : { opacity: 0, y: 2 }; + const fadeInitial = prefersReducedMotion + ? { opacity: 1 } + : { opacity: 0, y: 2 }; const fadeAnimate = { opacity: 1, y: 0 }; - const fadeExit = prefersReducedMotion ? { opacity: 1 } : { opacity: 0, y: -2 }; + const fadeExit = prefersReducedMotion + ? { opacity: 1 } + : { opacity: 0, y: -2 }; useEffect(() => { setShell(defaultShell); @@ -133,9 +181,9 @@ export function UpdateStudioInstructions({ onClick={() => setShell("unix")} className={cn( "px-0.5 py-0.5 font-medium transition-colors", - !windows - ? "text-foreground" - : "text-muted-foreground hover:text-emerald-600", + windows + ? "text-muted-foreground hover:text-emerald-600" + : "text-foreground", )} aria-pressed={!windows} > @@ -143,43 +191,157 @@ export function UpdateStudioInstructions({
- - - {getStudioUpdateInstructionLine(shell)} - - - -

- If that fails or unsloth studio update is unavailable, run: -

- - + {loadingInstallSource ? ( +

+ Checking how Studio was installed… +

+ ) : localInstallSource ? ( + <> +

+ Source or local install detected. To avoid replacing it with PyPI, + update from the checkout or source you originally installed from. +

+ {checkoutInstallSource ? ( + <> +

+ Pull latest changes from your Unsloth repo checkout, then update + Studio locally: +

+ + +

+ If the Studio update command is unavailable, run the local + installer from that checkout: +

+ + + + + + + ) : null} + {packagedSourceInstall ? ( + <> +

+ This looks like a source or VCS package install. Reinstall from + the original local path or Git URL you used. +

+

+ If you still have the Unsloth repo checkout, run the local + installer from that checkout: +

+ + + + + + + ) : null} +

+ Restart Studio after updating for changes to take effect. +

+ + ) : unknownInstallSource ? ( + <> +

+ Studio could not detect how it was installed. Check how you + installed Studio first, then choose the matching update path. +

+

+ For curl or PyPI installs, run: +

- - -

- Restart Studio after updating for changes to take effect. -

+

+ For local checkout installs, update from that checkout instead and + use the local update command: +

+ +

+ Restart Studio after updating for changes to take effect. +

+ + ) : ( + <> + + + {getStudioUpdateInstructionLine(shell)} + + + +

+ If that fails or unsloth studio update is unavailable, run: +

+ + + + + +

+ Restart Studio after updating for changes to take effect. +

+ + )} ); } diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx index 9629900237..68d82a804a 100644 --- a/studio/frontend/src/features/settings/tabs/about-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx @@ -1,12 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Button } from "@/components/ui/button"; import { ShutdownDialog } from "@/components/shutdown-dialog"; -import { UpdateStudioInstructions } from "../components/update-studio-instructions"; +import { Button } from "@/components/ui/button"; import { usePlatformStore } from "@/config/env"; -import { apiUrl } from "@/lib/api-base"; -import { removeTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard"; +import { getAuthToken } from "@/features/auth"; +import { removeTrainingUnloadGuard } from "@/features/training"; +import { apiUrl, isTauri } from "@/lib/api-base"; import { ArrowUpRight01Icon, Book03Icon, @@ -18,28 +18,108 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useState } from "react"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; +import { + type UpdateInstallSource, + UpdateStudioInstructions, +} from "../components/update-studio-instructions"; + +type ApiObject = Record; + +const INSTALL_SOURCE_KEY = "install_source"; + +const UPDATE_INSTALL_SOURCES = new Set([ + "pypi", + "editable", + "local_path", + "vcs", + "local_repo", + "unknown", +]); + +function isUpdateInstallSource(value: unknown): value is UpdateInstallSource { + return ( + typeof value === "string" && + UPDATE_INSTALL_SOURCES.has(value as UpdateInstallSource) + ); +} + +async function fetchStudioVersions(): Promise<{ + packageVersion: string | null; + studioVersion: string | null; +}> { + try { + const res = await fetch(apiUrl("/api/health")); + if (!res.ok) { + return { packageVersion: null, studioVersion: null }; + } + const data = (await res.json()) as ApiObject; + const packageVersion = data.version; + const studioVersion = data.studio_version; + return { + packageVersion: + typeof packageVersion === "string" ? packageVersion : null, + studioVersion: typeof studioVersion === "string" ? studioVersion : null, + }; + } catch { + return { packageVersion: null, studioVersion: null }; + } +} + +async function fetchInstallSource(): Promise { + if (isTauri) { + return "unknown"; + } + + const token = getAuthToken(); + if (!token) { + return "unknown"; + } + + try { + const headers = new Headers(); + headers.set("Authorization", `Bearer ${token}`); + const res = await fetch(apiUrl("/api/studio/install-source"), { headers }); + if (!res.ok) { + return "unknown"; + } + const data = (await res.json()) as ApiObject; + const installSource = data[INSTALL_SOURCE_KEY]; + return isUpdateInstallSource(installSource) ? installSource : "unknown"; + } catch { + return "unknown"; + } +} export function AboutTab() { const deviceType = usePlatformStore((s) => s.deviceType); const defaultShell = deviceType === "windows" ? "windows" : "unix"; const [shutdownOpen, setShutdownOpen] = useState(false); - const [version, setVersion] = useState("dev"); + const [packageVersion, setPackageVersion] = useState("dev"); + const [studioVersion, setStudioVersion] = useState("dev"); + const [installSource, setInstallSource] = useState< + UpdateInstallSource | "loading" + >("loading"); useEffect(() => { let canceled = false; - (async () => { - try { - const res = await fetch(apiUrl("/api/health")); - if (!res.ok) return; - const data = (await res.json()) as { version?: string }; - if (!canceled && data.version) { - setVersion(data.version); - } - } catch { - // fall back to dev label + fetchStudioVersions().then((nextVersions) => { + if (canceled) { + return; } - })(); + if (nextVersions.packageVersion) { + setPackageVersion(nextVersions.packageVersion); + } + if (nextVersions.studioVersion) { + setStudioVersion(nextVersions.studioVersion); + } + }); + + fetchInstallSource().then((nextInstallSource) => { + if (!canceled) { + setInstallSource(nextInstallSource); + } + }); return () => { canceled = true; @@ -56,14 +136,25 @@ export function AboutTab() { - - {version} + + + {studioVersion} + + + + + {packageVersion} +
- +
@@ -99,7 +190,10 @@ export function AboutTab() { rel="noopener noreferrer" className="inline-flex items-center gap-1 text-xs font-medium text-muted-foreground hover:text-foreground" > - + Report an issue @@ -108,7 +202,7 @@ export function AboutTab() { diff --git a/studio/frontend/src/features/training/hooks/use-training-unload-guard.ts b/studio/frontend/src/features/training/hooks/use-training-unload-guard.ts index 78c1e3d5c3..1808cf5888 100644 --- a/studio/frontend/src/features/training/hooks/use-training-unload-guard.ts +++ b/studio/frontend/src/features/training/hooks/use-training-unload-guard.ts @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useEffect } from "react"; -import { useTrainingRuntimeStore } from "@/features/training"; +import { useTrainingRuntimeStore } from "../stores/training-runtime-store"; let currentHandler: ((e: BeforeUnloadEvent) => void) | null = null; @@ -13,14 +13,18 @@ let currentHandler: ((e: BeforeUnloadEvent) => void) | null = null; export function useTrainingUnloadGuard() { useEffect(() => { const handler = (e: BeforeUnloadEvent) => { - if (!useTrainingRuntimeStore.getState().isTrainingRunning) return; + if (!useTrainingRuntimeStore.getState().isTrainingRunning) { + return; + } e.preventDefault(); e.returnValue = ""; }; currentHandler = handler; window.addEventListener("beforeunload", handler); return () => { - if (currentHandler === handler) currentHandler = null; + if (currentHandler === handler) { + currentHandler = null; + } window.removeEventListener("beforeunload", handler); }; }, []); diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts index 571cfb6ff5..83d8edba75 100644 --- a/studio/frontend/src/features/training/index.ts +++ b/studio/frontend/src/features/training/index.ts @@ -16,7 +16,11 @@ export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-st export { uploadTrainingDataset } from "./api/datasets-api"; export { listLocalModels } from "./api/models-api"; export type { LocalModelInfo } from "./api/models-api"; -export type { TrainingPhase, TrainingViewData, TrainingSeriesPoint } from "./types/runtime"; +export type { + TrainingPhase, + TrainingViewData, + TrainingSeriesPoint, +} from "./types/runtime"; export type { TrainingRunSummary, TrainingRunListResponse, diff --git a/studio/frontend/src/hooks/use-web-update-check.ts b/studio/frontend/src/hooks/use-web-update-check.ts new file mode 100644 index 0000000000..1ad0fe544c --- /dev/null +++ b/studio/frontend/src/hooks/use-web-update-check.ts @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { getAuthToken } from "@/features/auth"; +import { apiUrl, isTauri } from "@/lib/api-base"; +import { useCallback, useEffect, useState } from "react"; + +const WEB_UPDATE_CHECK_DELAY_MS = 5000; +const DISMISS_PREFIX = "unsloth_web_update_dismissed"; +const CAN_SHOW_KEY = "can_show_web_notification"; +const UPDATE_AVAILABLE_KEY = "update_available"; +const INSTALL_SOURCE_KEY = "install_source"; +const LATEST_VERSION_KEY = "latest_version"; +const CURRENT_VERSION_KEY = "current_version"; +const CHECKED_AT_KEY = "checked_at"; + +type ApiObject = Record; + +export type WebUpdateInstallSource = + | "pypi" + | "editable" + | "local_path" + | "vcs" + | "local_repo" + | "unknown"; + +export interface WebUpdateStatus { + currentVersion: string; + latestVersion: string; + installSource: "pypi"; + checkedAt: string; +} + +interface UseWebUpdateCheckOptions { + enabled?: boolean; + delayMs?: number; +} + +function stringField(value: ApiObject, key: string): string | null { + const field = value[key]; + return typeof field === "string" ? field : null; +} + +function toDisplayableUpdateStatus(value: unknown): WebUpdateStatus | null { + if (!value || typeof value !== "object") { + return null; + } + + const status = value as ApiObject; + const latestVersion = stringField(status, LATEST_VERSION_KEY); + const currentVersion = stringField(status, CURRENT_VERSION_KEY); + const checkedAt = stringField(status, CHECKED_AT_KEY); + if ( + status[CAN_SHOW_KEY] !== true || + status[UPDATE_AVAILABLE_KEY] !== true || + status[INSTALL_SOURCE_KEY] !== "pypi" || + !latestVersion || + !currentVersion || + !checkedAt + ) { + return null; + } + + return { + currentVersion, + latestVersion, + installSource: "pypi", + checkedAt, + }; +} + +function dismissalKey(status: WebUpdateStatus): string { + return `${DISMISS_PREFIX}:${status.installSource}:${status.latestVersion}`; +} + +function isDismissed(status: WebUpdateStatus): boolean { + if (typeof window === "undefined") { + return true; + } + try { + return window.localStorage.getItem(dismissalKey(status)) !== null; + } catch { + return false; + } +} + +function markDismissed(status: WebUpdateStatus): void { + if (typeof window === "undefined") { + return; + } + try { + window.localStorage.setItem(dismissalKey(status), String(Date.now())); + } catch { + // Ignore storage failures; the banner can still be dismissed in-memory. + } +} + +async function fetchDisplayableUpdateStatus(): Promise { + const token = getAuthToken(); + if (!token) { + return null; + } + + const headers = new Headers(); + headers.set("Authorization", `Bearer ${token}`); + const res = await fetch(apiUrl("/api/studio/update-status"), { headers }); + if (!res.ok) { + return null; + } + + return toDisplayableUpdateStatus(await res.json()); +} + +export function useWebUpdateCheck({ + enabled = true, + delayMs = WEB_UPDATE_CHECK_DELAY_MS, +}: UseWebUpdateCheckOptions = {}) { + const [status, setStatus] = useState(null); + + useEffect(() => { + if (isTauri || !enabled || !getAuthToken()) { + const clearTimer = window.setTimeout(() => setStatus(null), 0); + return () => window.clearTimeout(clearTimer); + } + + let canceled = false; + const timer = window.setTimeout(() => { + fetchDisplayableUpdateStatus() + .then((nextStatus) => { + if (canceled) { + return; + } + setStatus(nextStatus && !isDismissed(nextStatus) ? nextStatus : null); + }) + .catch(() => { + if (!canceled) { + setStatus(null); + } + }); + }, delayMs); + + return () => { + canceled = true; + window.clearTimeout(timer); + }; + }, [delayMs, enabled]); + + const dismiss = useCallback(() => { + setStatus((current) => { + if (current) { + markDismissed(current); + } + return null; + }); + }, []); + + return { status: enabled && !isTauri ? status : null, dismiss }; +} diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 158a22ebd0..7c933bd612 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -205,8 +205,12 @@ class AssetChoice: name: str url: str source_label: str + # Paired runtime archive (Windows CUDA cudart bundle). When set, + # install_from_archives also downloads it and overlays its DLLs on + # top of the main install. See unslothai/unsloth#5106. runtime_name: str | None = None runtime_url: str | None = None + runtime_sha256: str | None = None is_ready_bundle: bool = False install_kind: str = "" bundle_profile: str | None = None @@ -2922,6 +2926,30 @@ def windows_cuda_attempts( + ",".join(windows_cuda_upstream_asset_names(llama_tag, runtime)) ) continue + # Pair the cudart bundle when upstream ships it. Without this + # the binary needs a system CUDA toolkit on PATH at runtime + # (#5106). Only pair when the selected main archive is the + # binary archive, not the cudart archive itself. + runtime_archive_name: str | None = None + runtime_archive_url: str | None = None + if selected_name.startswith(f"llama-"): + cudart_name = f"cudart-llama-bin-win-cuda-{runtime}-x64.zip" + cudart_url = upstream_assets.get(cudart_name) + if cudart_url and cudart_url != asset_url: + runtime_archive_name = cudart_name + runtime_archive_url = cudart_url + attempt_log = list(selection_log) + [ + f"windows_cuda_selection: selected {selected_name} runtime={runtime}" + ] + if runtime_archive_name: + attempt_log.append( + f"windows_cuda_selection: paired runtime archive {runtime_archive_name}" + ) + else: + attempt_log.append( + "windows_cuda_selection: no paired runtime archive found; " + "binary will rely on a system CUDA toolkit at runtime" + ) attempts.append( AssetChoice( repo = UPSTREAM_REPO, @@ -2931,10 +2959,9 @@ def windows_cuda_attempts( source_label = "upstream", install_kind = "windows-cuda", runtime_line = runtime_line, - selection_log = list(selection_log) - + [ - f"windows_cuda_selection: selected {selected_name} runtime={runtime}" - ], + runtime_name = runtime_archive_name, + runtime_url = runtime_archive_url, + selection_log = attempt_log, ) ) return attempts @@ -2982,6 +3009,24 @@ def published_windows_cuda_attempts( asset_url = release.assets.get(artifact.asset_name) if not asset_url: continue + # See windows_cuda_attempts: pair the cudart bundle. + runtime_archive_name: str | None = None + runtime_archive_url: str | None = None + if artifact.asset_name.startswith("llama-"): + runtime = runtime_by_line[runtime_line] + cudart_name = f"cudart-llama-bin-win-cuda-{runtime}-x64.zip" + cudart_url = release.assets.get(cudart_name) + if cudart_url and cudart_url != asset_url: + runtime_archive_name = cudart_name + runtime_archive_url = cudart_url + attempt_log = list(ordered_attempt.selection_log or []) + [ + "windows_cuda_selection: selected published asset " + f"{artifact.asset_name} for runtime_line={runtime_line}" + ] + if runtime_archive_name: + attempt_log.append( + f"windows_cuda_selection: paired published runtime archive {runtime_archive_name}" + ) attempts.append( AssetChoice( repo = release.repo, @@ -2991,11 +3036,9 @@ def published_windows_cuda_attempts( source_label = "published", install_kind = "windows-cuda", runtime_line = runtime_line, - selection_log = list(ordered_attempt.selection_log or []) - + [ - "windows_cuda_selection: selected published asset " - f"{artifact.asset_name} for runtime_line={runtime_line}" - ], + runtime_name = runtime_archive_name, + runtime_url = runtime_archive_url, + selection_log = attempt_log, ) ) break @@ -3701,6 +3744,17 @@ def overlay_directory_for_choice( return path +def paired_runtime_dll_patterns(choice: AssetChoice) -> list[str]: + """Filename patterns the paired runtime archive is allowed to drop + into the install. Used for the second copy_globs pass in + install_from_archives, narrower than runtime_patterns_for_choice so + the runtime archive cannot overwrite main-archive payload like + llama-server.exe. Only Windows CUDA has paired runtimes today.""" + if choice.install_kind == "windows-cuda": + return ["cudart64_*.dll", "cublas64_*.dll", "cublasLt64_*.dll"] + return [] + + def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: if choice.install_kind in {"linux-cpu", "linux-cuda", "linux-rocm"}: return [ @@ -4020,14 +4074,52 @@ def install_from_archives( install_dir.mkdir(parents = True, exist_ok = True) extract_dir = Path(tempfile.mkdtemp(prefix = "extract-", dir = work_dir)) + runtime_extract_dir: Path | None = None try: extract_archive(main_archive, extract_dir) + # Download the paired runtime archive into its own temp dir to + # avoid copy_globs's ambiguous-layout guard on shared names + # like LICENSE.txt. Two passes of copy_globs land both archives + # in the same overlay dir. Fixes #5106. + if choice.runtime_url and choice.runtime_name: + runtime_archive = work_dir / choice.runtime_name + log( + f"downloading paired runtime archive {choice.runtime_name} " + f"from {choice.source_label} release" + ) + download_file_verified( + choice.runtime_url, + runtime_archive, + expected_sha256 = choice.runtime_sha256, + label = f"prebuilt runtime archive {choice.runtime_name}", + ) + runtime_extract_dir = Path( + tempfile.mkdtemp(prefix = "extract-runtime-", dir = work_dir) + ) + extract_archive(runtime_archive, runtime_extract_dir) source_dir = extract_dir overlay_dir = overlay_directory_for_choice(install_dir, choice, host) copy_globs( source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True ) + if runtime_extract_dir is not None: + # The runtime archive only contributes the CUDA DLLs. + # Restrict the overlay to the cudart bundle's known + # filenames (cudart64_X.dll / cublas64_X.dll / + # cublasLt64_X.dll) rather than the broad ``*.exe`` / + # ``*.dll`` set from runtime_patterns_for_choice, so a + # malformed runtime archive can never overwrite + # llama-server.exe or other main-archive payload. The + # upstream cudart-llama-bin-win-cuda-X.Y-x64.zip currently + # ships exactly these three DLLs (verified against b9103 + # cuda-12.4 and cuda-13.1 bundles). + copy_globs( + runtime_extract_dir, + overlay_dir, + paired_runtime_dll_patterns(choice), + required = False, + ) copy_globs( source_dir, install_dir, @@ -4036,6 +4128,8 @@ def install_from_archives( ) finally: remove_tree(extract_dir) + if runtime_extract_dir is not None: + remove_tree(runtime_extract_dir) if host.is_windows: exec_dir = install_dir / "build" / "bin" / "Release" @@ -4239,8 +4333,27 @@ def python_runtime_dirs() -> list[str]: for root in search_roots: if not root.is_dir(): continue + # ``nvidia//lib`` -- Linux convention; harmless on Windows + # where the directory simply does not exist on real wheels. candidates.extend(root.glob("nvidia/*/lib")) + # ``nvidia//bin`` -- legacy modular Windows wheels + # (``nvidia-cuda-runtime-cu12``, ``nvidia-cublas-cu12``). candidates.extend(root.glob("nvidia/*/bin")) + # ``nvidia//bin/x86_64`` and ``.../bin/x64`` -- current + # CUDA 13 Windows wheel layout (the unsuffixed + # ``nvidia-cuda-runtime`` 13.x and ``nvidia-cublas`` 13.x + # packages ship under ``nvidia/cu13/bin/x86_64/cudart64_13.dll``). + # Without these, Windows preflight CUDA detection misses cu13 + # installs and falls back to the upstream cudart bundle path + # even when usable DLLs are already on disk (#5106). Kept in + # sync with the backend resolver + # ``llama_cpp.LlamaCppBackend._windows_pip_nvidia_dll_dirs``. + candidates.extend(root.glob("nvidia/*/bin/x86_64")) + candidates.extend(root.glob("nvidia/*/bin/x64")) + # ``nvidia//Library/bin`` -- conda-style wheel repacks. + candidates.extend(root.glob("nvidia/*/Library/bin")) + candidates.extend(root.glob("nvidia/*/Library/bin/x86_64")) + candidates.extend(root.glob("nvidia/*/Library/bin/x64")) candidates.extend(root.glob("torch/lib")) return dedupe_existing_dirs(candidates) @@ -4743,6 +4856,17 @@ def apply_approved_hashes( missing_assets.append(attempt.name) continue attempt.expected_sha256 = approved.sha256 + # Resolve the paired runtime archive's hash too. Drop the pair + # if the manifest does not list it -- never install an + # unverified archive. + if attempt.runtime_name and attempt.runtime_url: + runtime_approved = checksums.artifacts.get(attempt.runtime_name) + if runtime_approved is None: + attempt.runtime_name = None + attempt.runtime_url = None + attempt.runtime_sha256 = None + else: + attempt.runtime_sha256 = runtime_approved.sha256 approved_attempts.append(attempt) if not approved_attempts: missing_text = ", ".join(missing_assets) if missing_assets else "none" @@ -4906,24 +5030,19 @@ def write_prebuilt_metadata( approved_checksums, llama_tag, ) - fingerprint_payload = { - "published_repo": approved_checksums.repo, - "release_tag": release_tag, - "upstream_tag": llama_tag, - "asset": choice.name, - "asset_sha256": choice.expected_sha256, - "source": choice.source_label, - "source_asset": source_asset_name, - "source_sha256": source_sha256, - "runtime_line": choice.runtime_line, - "bundle_profile": choice.bundle_profile, - "coverage_class": choice.coverage_class, - } - fingerprint = hashlib.sha256( - json.dumps(fingerprint_payload, sort_keys = True, separators = (",", ":")).encode( - "utf-8" - ) - ).hexdigest() + # expected_install_fingerprint is the source of truth for what the + # fingerprint must contain. Calling it here -- instead of inlining a + # parallel payload -- prevents drift where new keys (e.g. the cudart + # pair fields added for #5106) are added to one side but not the + # other, which would cause every install to look stale. + fingerprint = expected_install_fingerprint( + llama_tag = llama_tag, + release_tag = release_tag, + choice = choice, + approved_checksums = approved_checksums, + ) + if fingerprint is None: + raise PrebuiltFallback(f"cannot compute install fingerprint for {choice.name}") metadata = { "requested_tag": requested_tag, "tag": llama_tag, @@ -4974,6 +5093,14 @@ def expected_install_fingerprint( "source_asset": source_asset_name, "source_sha256": source_sha256, "runtime_line": choice.runtime_line, + # Including the paired runtime archive (Windows cudart bundle) + # in the fingerprint is what forces existing #5106 installs to + # refresh: pre-PR installs hashed nothing in this slot, post-PR + # paired installs hash the cudart sha. Without these two keys + # an existing cudart-less install would keep matching the new + # choice and never re-overlay the cudart DLLs. + "runtime_asset": choice.runtime_name, + "runtime_sha256": choice.runtime_sha256, "bundle_profile": choice.bundle_profile, "coverage_class": choice.coverage_class, } @@ -5034,7 +5161,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: if choice.install_kind == "windows-cpu": return [["llama.dll"]] if choice.install_kind == "windows-cuda": - return [["llama.dll"], ["ggml-cuda.dll"]] + groups = [["llama.dll"], ["ggml-cuda.dll"]] + # When the cudart bundle was paired in (#5106) require all + # three of its DLLs alongside the main archive's payload. + # install_kind alone is not enough -- legacy installs without + # the cudart pair must still pass the health check on the + # no-pair fallback path, otherwise pair-less builds would loop + # on reinstall forever. The upstream cudart bundle ships + # cudart64_X.dll + cublas64_X.dll + cublasLt64_X.dll; missing + # any one of them still breaks GPU initialisation. + if choice.runtime_name: + groups.append(["cudart64_*.dll"]) + groups.append(["cublas64_*.dll"]) + groups.append(["cublasLt64_*.dll"]) + return groups if choice.install_kind == "windows-hip": return [["llama.dll"], ["*hip*.dll"]] return [] diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py index 928fa242eb..3deeb38cda 100644 --- a/tests/studio/_playwright_robust.py +++ b/tests/studio/_playwright_robust.py @@ -21,7 +21,9 @@ It does NOT depend on pytest -- both consumers run as plain Python. from __future__ import annotations import json +import os import sys +import threading import time import urllib.error import urllib.request @@ -404,3 +406,142 @@ def dump_diagnostics( except Exception as exc: if info is not None: info(f"diagnostics: json sidecar {name} failed: {exc}") + + +# ───────────────────────────────────────────────────────────────────── +# Bounded in-page fetch. +# ───────────────────────────────────────────────────────────────────── +# +# Playwright's `page.evaluate(...)` has no `timeout=` argument. If the +# JS body awaits a fetch that never resolves (the renderer's network +# thread wedges, the server accepts the connection but never replies, +# the macos-14 free runner under --single-process Chromium loses its +# IPC pipe), the entire Python script hangs until the runner-level +# timeout fires. Run 25696797934 / job 75446949358 on PR #5387 showed +# this exact failure: studio.log went idle after the chat surface +# mounted, no further requests reached the server, and Playwright +# burned 27+ minutes on a single page.evaluate(fetch /api/inference/ +# load) before the 30-min runner cancel. +# +# `evaluate_fetch` wraps the fetch in an AbortController.signal so the +# JS side resolves either with a real response or with a synthetic +# `{status: 0, error: "AbortError..."}` after `timeout_ms` ms. Either +# way page.evaluate returns and the script proceeds (or fails) with +# a debuggable signal instead of a silent wedge. +def evaluate_fetch( + page: Any, + url: str, + *, + method: str = "GET", + headers: dict[str, str] | None = None, + body: Any = None, + timeout_ms: int = 20_000, +) -> dict[str, Any]: + """Run `fetch(url, opts)` inside the page with an AbortSignal deadline. + + Returns `{"status": int, "body": parsed_or_text, "error": str|None}`. + On AbortSignal timeout returns `{"status": 0, "body": None, "error": + "AbortError: ..."}`. Callers should treat `status == 0` (or any + non-None `error`) as a transport failure rather than an HTTP + response. + + `body` may be a `str` (sent verbatim) or a `dict`/`list` (JSON- + encoded here). Pass headers explicitly when you need + `Content-Type: application/json` or an `Authorization` bearer. + """ + body_arg: str | None + if body is None: + body_arg = None + elif isinstance(body, (str, bytes)): + body_arg = body if isinstance(body, str) else body.decode("utf-8") + else: + body_arg = json.dumps(body) + js = """ + async ({url, method, headers, body, timeoutMs}) => { + const ctrl = new AbortController(); + const t = setTimeout(() => ctrl.abort(), timeoutMs); + try { + const opts = {method: method, headers: headers, signal: ctrl.signal}; + if (body !== null) opts.body = body; + const r = await fetch(url, opts); + clearTimeout(t); + let parsed; + try { + parsed = await r.json(); + } catch (_e) { + try { + parsed = await r.text(); + } catch (_e2) { + parsed = null; + } + } + return {status: r.status, body: parsed, error: null}; + } catch (e) { + clearTimeout(t); + return {status: 0, body: null, error: String(e)}; + } + } + """ + return page.evaluate( + js, + { + "url": url, + "method": method, + "headers": headers or {}, + "body": body_arg, + "timeoutMs": int(timeout_ms), + }, + ) + + +# ───────────────────────────────────────────────────────────────────── +# Wall-clock watchdog. +# ───────────────────────────────────────────────────────────────────── +# +# Even with every action and fetch bounded, a sufficiently strange +# wedge inside the browser (a CPU-pinned JS infinite loop, a renderer +# crash that doesn't propagate to Playwright, an asyncio deadlock in +# the sync wrapper) can still hang the script. The watchdog is a +# daemon Timer that calls `os._exit(2)` after `deadline_s` seconds, +# printing the wedge location to stderr so the CI log shows where the +# script was at force-kill time. The exit code matches "test failure +# by deadline" so the workflow's `set -e` propagates correctly. +# +# Pick `deadline_s` generously enough to cover the slowest healthy +# run -- macos-14 free runners with cold caches measure ~7-9 min for +# the comprehensive chat UI test. 12 minutes (720 s) leaves headroom +# without amplifying every real wedge to the 30-min runner-level cap. +def install_wall_clock_watchdog( + deadline_s: float, + *, + label: str = "playwright", + info: Callable[[str], None] | None = None, +) -> threading.Timer: + """Start a daemon Timer that hard-exits the process at `deadline_s`. + + Returns the Timer so the caller can `.cancel()` it on clean exit. + The Timer is daemonised; if the script exits normally before the + deadline the Timer dies with the process even without an explicit + cancel. + """ + + def _kaboom() -> None: + msg = ( + f"[{label}] WATCHDOG: hit {deadline_s:.0f}s wall-clock " + f"deadline; forcing exit(2). The script wedged somewhere " + f"the per-action timeouts could not bound. Inspect the " + f"most recent step printed above to localise." + ) + try: + sys.stderr.write(msg + "\n") + sys.stderr.flush() + except Exception: + pass + os._exit(2) + + timer = threading.Timer(deadline_s, _kaboom) + timer.daemon = True + timer.start() + if info is not None: + info(f"watchdog armed: hard-exit at {deadline_s:.0f}s") + return timer diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 128b90cfe2..b44554504f 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -769,7 +769,11 @@ def write_linux_install_shape(install_dir: Path) -> None: def write_windows_install_shape( - install_dir: Path, *, include_llama_dll: bool = True, include_cuda_dll: bool = False + install_dir: Path, + *, + include_llama_dll: bool = True, + include_cuda_dll: bool = False, + include_cudart_dlls: bool = False, ) -> None: runtime_dir = install_dir / "build" / "bin" / "Release" runtime_dir.mkdir(parents = True, exist_ok = True) @@ -779,6 +783,11 @@ def write_windows_install_shape( (runtime_dir / "llama.dll").write_bytes(b"DLL") if include_cuda_dll: (runtime_dir / "ggml-cuda.dll").write_bytes(b"DLL") + if include_cudart_dlls: + # cudart bundle DLLs that ship in cudart-llama-bin-win-cuda-*-x64.zip + (runtime_dir / "cudart64_12.dll").write_bytes(b"DLL") + (runtime_dir / "cublas64_12.dll").write_bytes(b"DLL") + (runtime_dir / "cublasLt64_12.dll").write_bytes(b"DLL") (install_dir / "convert_hf_to_gguf.py").write_text( "#!/usr/bin/env python3\n", encoding = "utf-8" ) @@ -1153,6 +1162,330 @@ def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path: assert existing_install_matches_plan(install_dir, host, plan) is False +def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart( + tmp_path: Path, +): + """When the choice ships a paired cudart bundle (#5106), the install + is considered stale unless cudart64_*.dll and cublas64_*.dll are + actually on disk. Otherwise existing broken installs would keep + matching and skip the reinstall that drops cudart in.""" + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_windows_install_shape( + install_dir, + include_llama_dll = True, + include_cuda_dll = True, + include_cudart_dlls = True, + ) + + host = HostInfo( + system = "Windows", + machine = "AMD64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = (12, 4), + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = True, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-win-cuda-12.4-x64.zip", + url = "https://example.com/x.zip", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = "a" * 64, + runtime_name = "cudart-llama-bin-win-cuda-12.4-x64.zip", + runtime_url = "https://example.com/cudart.zip", + runtime_sha256 = "c" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + choice.runtime_name: ApprovedArtifactHash( + asset_name = choice.runtime_name, + sha256 = choice.runtime_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + # Fully populated install (main archive + cudart DLLs) matches. + assert existing_install_matches_plan(install_dir, host, plan) is True + + # cublas missing -- stale, must reinstall. + (install_dir / "build" / "bin" / "Release" / "cublas64_12.dll").unlink() + assert existing_install_matches_plan(install_dir, host, plan) is False + + # cudart missing -- stale, must reinstall. + write_windows_install_shape( + install_dir, + include_llama_dll = True, + include_cuda_dll = True, + include_cudart_dlls = True, + ) + (install_dir / "build" / "bin" / "Release" / "cudart64_12.dll").unlink() + assert existing_install_matches_plan(install_dir, host, plan) is False + + # cublasLt missing -- stale, must reinstall. The upstream cudart + # bundle ships all three of cudart / cublas / cublasLt; a user with + # cudart + cublas but no cublasLt is still missing a required GPU + # initialisation DLL and Studio must refresh the install. + write_windows_install_shape( + install_dir, + include_llama_dll = True, + include_cuda_dll = True, + include_cudart_dlls = True, + ) + (install_dir / "build" / "bin" / "Release" / "cublasLt64_12.dll").unlink() + assert existing_install_matches_plan(install_dir, host, plan) is False + + +def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check( + tmp_path: Path, +): + """If the choice has no paired runtime archive (manifest dropped it, + or upstream did not ship cudart), legacy installs without cudart on + disk must still pass the health check -- otherwise the installer + would loop on reinstall forever because install_from_archives has no + cudart source to drop in.""" + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_windows_install_shape( + install_dir, + include_llama_dll = True, + include_cuda_dll = True, + include_cudart_dlls = False, + ) + + host = HostInfo( + system = "Windows", + machine = "AMD64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = (12, 4), + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = True, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-win-cuda-12.4-x64.zip", + url = "https://example.com/x.zip", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is True + + +def test_existing_install_fingerprint_changes_when_cudart_pair_added( + tmp_path: Path, +): + """Existing pre-#5322 Windows CUDA installs (no paired cudart) must + be treated as stale once the choice gains a runtime archive, + otherwise the fingerprint match would keep skipping the reinstall + that drops the cudart DLLs in. This is the install-cache half of the + #5106 fix -- the health-check half lives in the test above.""" + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_windows_install_shape( + install_dir, + include_llama_dll = True, + include_cuda_dll = True, + include_cudart_dlls = False, + ) + + host = HostInfo( + system = "Windows", + machine = "AMD64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = (12, 4), + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = True, + ) + legacy_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-win-cuda-12.4-x64.zip", + url = "https://example.com/x.zip", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = "a" * 64, + ) + paired_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-win-cuda-12.4-x64.zip", + url = "https://example.com/x.zip", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = "a" * 64, + runtime_name = "cudart-llama-bin-win-cuda-12.4-x64.zip", + runtime_url = "https://example.com/cudart.zip", + runtime_sha256 = "c" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + legacy_choice.name: ApprovedArtifactHash( + asset_name = legacy_choice.name, + sha256 = legacy_choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + paired_choice.runtime_name: ApprovedArtifactHash( + asset_name = paired_choice.runtime_name, + sha256 = paired_choice.runtime_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + }, + ) + + # Install metadata was written for the legacy (no-pair) choice. + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = legacy_choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + # New plan offers the paired choice -- fingerprint must differ so + # the install is refreshed. The health check would also catch this + # because cudart64_*.dll is missing on disk; we test the fingerprint + # half explicitly by comparing the two fingerprints directly. + legacy_fingerprint = INSTALL_LLAMA_PREBUILT.expected_install_fingerprint( + llama_tag = "b9001", + release_tag = "release-1", + choice = legacy_choice, + approved_checksums = checksums, + ) + paired_fingerprint = INSTALL_LLAMA_PREBUILT.expected_install_fingerprint( + llama_tag = "b9001", + release_tag = "release-1", + choice = paired_choice, + approved_checksums = checksums, + ) + assert legacy_fingerprint != paired_fingerprint, ( + "expected_install_fingerprint must hash runtime_name/runtime_sha256 " + "so pre-#5322 installs are not falsely considered up-to-date" + ) + + paired_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [paired_choice], + approved_checksums = checksums, + ) + assert existing_install_matches_plan(install_dir, host, paired_plan) is False + + def test_existing_install_matches_plan_macos_requires_dylibs(tmp_path: Path): install_dir = tmp_path / "llama.cpp" install_dir.mkdir() @@ -2050,3 +2383,184 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete_maco ) is False ) + + +def test_paired_runtime_dll_patterns_excludes_executables() -> None: + """The paired runtime archive must only contribute CUDA DLLs to + the install. The narrow pattern list -- not the broad + runtime_patterns_for_choice ``*.exe`` / ``*.dll`` -- is what + prevents a malformed cudart bundle from overwriting + llama-server.exe at install time. + """ + paired_runtime_dll_patterns = INSTALL_LLAMA_PREBUILT.paired_runtime_dll_patterns + paired_choice = AssetChoice( + repo = "x", + tag = "t", + name = "llama-b9001-bin-win-cuda-12.4-x64.zip", + url = "u", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = "a" * 64, + runtime_name = "cudart-llama-bin-win-cuda-12.4-x64.zip", + runtime_url = "https://example.com/cudart.zip", + runtime_sha256 = "c" * 64, + ) + patterns = paired_runtime_dll_patterns(paired_choice) + assert "cudart64_*.dll" in patterns + assert "cublas64_*.dll" in patterns + assert "cublasLt64_*.dll" in patterns + assert "*.exe" not in patterns + assert "*.dll" not in patterns + + for kind in ( + "linux-cpu", + "linux-cuda", + "linux-rocm", + "macos-arm64", + "macos-x64", + "windows-cpu", + "windows-hip", + ): + non_windows = AssetChoice( + repo = "x", + tag = "t", + name = "x", + url = "u", + source_label = "published", + install_kind = kind, + expected_sha256 = "a" * 64, + ) + assert paired_runtime_dll_patterns(non_windows) == [] + + +def test_runtime_overlay_cannot_overwrite_main_archive_payload( + tmp_path: Path, +) -> None: + """End-to-end: a malformed runtime archive containing + ``llama-server.exe`` alongside the real cudart DLLs must NOT + replace the main archive's ``llama-server.exe``. + """ + install_from_archives = INSTALL_LLAMA_PREBUILT.install_from_archives + + work = tmp_path / "work" + install = tmp_path / "install" + archives = tmp_path / "archives" + work.mkdir() + install.mkdir() + archives.mkdir() + + main_zip = archives / "llama-b9001-bin-win-cuda-12.4-x64.zip" + runtime_zip = archives / "cudart-llama-bin-win-cuda-12.4-x64.zip" + with zipfile.ZipFile(main_zip, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr("llama-server.exe", b"MAIN-SERVER") + zf.writestr("llama-quantize.exe", b"MAIN-Q") + zf.writestr("llama.dll", b"DLL-llama") + zf.writestr("ggml-cuda.dll", b"DLL-ggml") + import hashlib + + main_sha = hashlib.sha256(main_zip.read_bytes()).hexdigest() + with zipfile.ZipFile(runtime_zip, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr("cudart64_12.dll", b"DLL-cudart") + zf.writestr("cublas64_12.dll", b"DLL-cublas") + zf.writestr("cublasLt64_12.dll", b"DLL-cublasLt") + zf.writestr("llama-server.exe", b"RUNTIME-OVERWRITE") + runtime_sha = hashlib.sha256(runtime_zip.read_bytes()).hexdigest() + + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = main_zip.name, + url = f"https://example.com/{main_zip.name}", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = main_sha, + runtime_name = runtime_zip.name, + runtime_url = f"https://example.com/{runtime_zip.name}", + runtime_sha256 = runtime_sha, + ) + host = HostInfo( + system = "Windows", + machine = "AMD64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = (12, 4), + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = True, + ) + + import shutil as _shutil + + orig_download = INSTALL_LLAMA_PREBUILT.download_file_verified + + def fake_download(url, target_path, *, expected_sha256 = None, label = None, **kw): + src = main_zip if "cudart" not in url else runtime_zip + _shutil.copy2(src, target_path) + if expected_sha256: + actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest() + if actual != expected_sha256: + raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback( + f"sha256 mismatch on {label}" + ) + + INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download + try: + install_from_archives(choice, host, install, work) + finally: + INSTALL_LLAMA_PREBUILT.download_file_verified = orig_download + + release_dir = install / "build" / "bin" / "Release" + server = release_dir / "llama-server.exe" + assert server.exists() + assert server.read_bytes() == b"MAIN-SERVER", ( + "runtime archive overwrote main llama-server.exe; " + f"got {server.read_bytes()!r}" + ) + for name in ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"): + assert (release_dir / name).exists(), f"missing {name}" + + +def test_python_runtime_dirs_covers_cu13_and_library_bin( + monkeypatch, tmp_path: Path +) -> None: + """Installer-side runtime DLL discovery must scan the same path + set as the backend ``_windows_pip_nvidia_dll_dirs``: legacy + ``nvidia//bin``, current ``nvidia//bin/x86_64`` + (cu13 layout), conda-style ``nvidia//Library/bin``, plus + ``torch/lib``. Otherwise installer preflight and backend launch + can disagree about which DLLs are actually present. + """ + import site as _site + + python_runtime_dirs = INSTALL_LLAMA_PREBUILT.python_runtime_dirs + + site_dir = tmp_path / "Lib" / "site-packages" + # cu12-style modular wheel + cu12_bin = site_dir / "nvidia" / "cuda_runtime" / "bin" + cu12_bin.mkdir(parents = True) + # cu13-style unsuffixed wheel + cu13_arch = site_dir / "nvidia" / "cu13" / "bin" / "x86_64" + cu13_arch.mkdir(parents = True) + # conda-style repack + library_bin = site_dir / "nvidia" / "cublas" / "Library" / "bin" + library_bin.mkdir(parents = True) + # PyTorch bundled-CUDA wheel + torch_lib = site_dir / "torch" / "lib" + torch_lib.mkdir(parents = True) + + monkeypatch.setattr(sys, "path", [str(site_dir)]) + monkeypatch.setattr(_site, "getsitepackages", lambda: [str(site_dir)]) + monkeypatch.setattr(_site, "getusersitepackages", lambda: "") + + dirs = python_runtime_dirs() + assert str(cu12_bin) in dirs + assert str(cu13_arch) in dirs + assert str(library_bin) in dirs + assert str(torch_lib) in dirs diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 5150d2d2ab..713af51e19 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -1839,6 +1839,126 @@ class TestWindowsCudaAttempts: assert result[0].name == "cudart-llama-bin-win-cuda-13.1-x64.zip" assert result[1].name == "cudart-llama-bin-win-cuda-12.4-x64.zip" + def test_cudart_runtime_archive_is_paired(self, monkeypatch): + # #5106: cudart bundle must surface on runtime_url so + # install_from_archives downloads it. + mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"]) + host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1)) + assets = { + f"llama-{self.TAG}-bin-win-cuda-13.1-x64.zip": f"https://example.com/llama-{self.TAG}-bin-win-cuda-13.1-x64.zip", + "cudart-llama-bin-win-cuda-13.1-x64.zip": "https://example.com/cudart-llama-bin-win-cuda-13.1-x64.zip", + f"llama-{self.TAG}-bin-win-cuda-12.4-x64.zip": f"https://example.com/llama-{self.TAG}-bin-win-cuda-12.4-x64.zip", + "cudart-llama-bin-win-cuda-12.4-x64.zip": "https://example.com/cudart-llama-bin-win-cuda-12.4-x64.zip", + } + result = windows_cuda_attempts(host, self.TAG, assets, None) + assert len(result) == 2 + # cuda13 first (host driver supports 13.1) + assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-13.1-x64.zip" + assert result[0].runtime_name == "cudart-llama-bin-win-cuda-13.1-x64.zip" + assert result[0].runtime_url == ( + "https://example.com/cudart-llama-bin-win-cuda-13.1-x64.zip" + ) + # cuda12 second + assert result[1].name == f"llama-{self.TAG}-bin-win-cuda-12.4-x64.zip" + assert result[1].runtime_name == "cudart-llama-bin-win-cuda-12.4-x64.zip" + + def test_no_runtime_archive_when_cudart_absent(self, monkeypatch): + # Older releases without the cudart split must still install. + mock_windows_runtime(monkeypatch, ["cuda12"]) + host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (12, 4)) + assets = { + f"llama-{self.TAG}-bin-win-cuda-12.4-x64.zip": f"https://example.com/llama-{self.TAG}-bin-win-cuda-12.4-x64.zip", + } + result = windows_cuda_attempts(host, self.TAG, assets, None) + assert len(result) == 1 + assert result[0].runtime_url is None + assert result[0].runtime_name is None + + def test_cudart_only_assets_do_not_self_pair(self, monkeypatch): + # Legacy cudart-only naming path must not self-pair. + mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"]) + host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1)) + assets = self._upstream("13.1", "12.4", current_names = True) + result = windows_cuda_attempts(host, self.TAG, assets, None) + assert len(result) == 2 + for attempt in result: + assert attempt.runtime_url is None + assert attempt.runtime_name is None + + +# =========================================================================== +# N.1. apply_approved_hashes -- runtime archive checksum threading +# =========================================================================== + + +class TestApplyApprovedHashesRuntimePair: + """Runtime archive must inherit a manifest hash, or be dropped.""" + + TAG = "b8508" + + def _runtime_paired_attempt(self) -> AssetChoice: + return AssetChoice( + repo = "unslothai/llama.cpp", + tag = self.TAG, + name = f"llama-{self.TAG}-bin-win-cuda-13.1-x64.zip", + url = f"https://x/llama-{self.TAG}-bin-win-cuda-13.1-x64.zip", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda13", + runtime_name = "cudart-llama-bin-win-cuda-13.1-x64.zip", + runtime_url = "https://x/cudart-llama-bin-win-cuda-13.1-x64.zip", + ) + + def test_runtime_hash_threaded_when_present(self): + attempt = self._runtime_paired_attempt() + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = self.TAG, + upstream_tag = self.TAG, + artifacts = { + attempt.name: ApprovedArtifactHash( + asset_name = attempt.name, + sha256 = "0" * 64, + repo = "unslothai/llama.cpp", + kind = "windows-cuda", + ), + "cudart-llama-bin-win-cuda-13.1-x64.zip": ApprovedArtifactHash( + asset_name = "cudart-llama-bin-win-cuda-13.1-x64.zip", + sha256 = "1" * 64, + repo = "unslothai/llama.cpp", + kind = "windows-cuda", + ), + }, + ) + result = apply_approved_hashes([attempt], checksums) + assert len(result) == 1 + assert result[0].expected_sha256 == "0" * 64 + assert result[0].runtime_sha256 == "1" * 64 + assert result[0].runtime_name == "cudart-llama-bin-win-cuda-13.1-x64.zip" + + def test_runtime_pair_dropped_when_hash_missing(self): + # Drop the pair rather than install an unverified runtime. + attempt = self._runtime_paired_attempt() + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = self.TAG, + upstream_tag = self.TAG, + artifacts = { + attempt.name: ApprovedArtifactHash( + asset_name = attempt.name, + sha256 = "0" * 64, + repo = "unslothai/llama.cpp", + kind = "windows-cuda", + ), + }, + ) + result = apply_approved_hashes([attempt], checksums) + assert len(result) == 1 + assert result[0].expected_sha256 == "0" * 64 + assert result[0].runtime_url is None + assert result[0].runtime_name is None + assert result[0].runtime_sha256 is None + # =========================================================================== # O. resolve_upstream_asset_choice -- platform routing diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 3f4ee6704c..aa1d38c4e1 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -55,7 +55,9 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) from _playwright_robust import ( # noqa: E402 chromium_launch_args, click_and_wait_for_response, + evaluate_fetch, install_view_transition_killer, + install_wall_clock_watchdog, is_benign_console_error, is_benign_page_error, recover_or_replace_page, @@ -85,6 +87,17 @@ STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1" # CI bump this without hard-coding a Mac branch in the test. TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000")) +# Wall-clock cap for the entire script. A healthy comprehensive run is +# 5-9 min; 12 min leaves headroom. Tunable via STUDIO_UI_WALL_TIMEOUT_S. +# See _playwright_robust.install_wall_clock_watchdog for rationale. +WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720")) + +# Per-fetch budget for in-page fetches. The /api/inference/load call is +# usually the slowest legitimate request: it pulls the model into the +# llama.cpp worker. Give it ~3 min on a cold cache, less elsewhere. +FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000")) +LOAD_FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_LOAD_TIMEOUT_MS", "180000")) + _n = [0] @@ -132,6 +145,11 @@ def parse_rgb(s): with sync_playwright() as p: + _watchdog = install_wall_clock_watchdog( + WALL_TIMEOUT_S, + label = "ui", + info = info, + ) # Pre-flight: bash-side wait_for already gated on /api/health # before launching us, but the macos-14 free runner has been # observed to surface a 200 /api/health while the auth DB is @@ -424,18 +442,18 @@ with sync_playwright() as p: "() => localStorage.getItem('unsloth_auth_refresh_token')", ) if refresh_token: - refresh = page.evaluate( - f"""async (rt) => {{ - const r = await fetch("{BASE}/api/auth/refresh", {{ - method: "POST", - headers: {{"Content-Type": "application/json"}}, - body: JSON.stringify({{refresh_token: rt}}), - }}); - return await r.json(); - }}""", - refresh_token, + refresh_resp = evaluate_fetch( + page, + f"{BASE}/api/auth/refresh", + method = "POST", + headers = {"Content-Type": "application/json"}, + body = {"refresh_token": refresh_token}, + timeout_ms = FETCH_TIMEOUT_MS, ) - token = refresh.get("access_token") + if refresh_resp.get("error"): + fail(f"/api/auth/refresh wedged: {refresh_resp['error']!r}") + refresh = refresh_resp.get("body") or {} + token = (refresh or {}).get("access_token") if not token: fail("could not obtain auth token after change-password") @@ -450,15 +468,18 @@ with sync_playwright() as p: "EXPECTED_DEFAULT_MODEL", "unsloth/gemma-4-E2B-it-GGUF", ) - defaults = page.evaluate( - f"""async (token) => {{ - const r = await fetch("{BASE}/api/models/list", {{ - headers: {{ "Authorization": "Bearer " + token }}, - }}); - return await r.json(); - }}""", - token, + defaults_resp = evaluate_fetch( + page, + f"{BASE}/api/models/list", + headers = {"Authorization": f"Bearer {token}"}, + timeout_ms = FETCH_TIMEOUT_MS, ) + if defaults_resp.get("error") or defaults_resp.get("status") != 200: + fail( + f"/api/models/list failed: status={defaults_resp.get('status')!r} " + f"error={defaults_resp.get('error')!r}" + ) + defaults = defaults_resp["body"] or {} if not defaults.get("default_models"): fail(f"/api/models/list returned no default_models: {defaults}") if defaults["default_models"][0] != EXPECTED_DEFAULT: @@ -479,8 +500,16 @@ with sync_playwright() as p: 'button:has-text("Qwen"), ' 'button:has-text("Llama")' ).first - if selector_btn.count() > 0: - sel_text = (selector_btn.text_content() or "").strip() + # Best-effort: the selector re-mounts as /api/models/list resolves, + # so use a short timeout and skip the snapshot on miss. + sel_text = "" + try: + sel_text = (selector_btn.text_content(timeout = 2_000) or "").strip() + except Exception as _sel_err: + info( + f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}" + ) + if sel_text: info(f"model selector button text: {sel_text!r}") shoot("03b-default-model-button") @@ -491,27 +520,35 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── step("load GGUF via /api/inference/load (uses session cookie)") # Token already fetched above; reuse it for the load call. - load_resp = page.evaluate(f"""async () => {{ - const r = await fetch("{BASE}/api/inference/load", {{ - method: "POST", - headers: {{ - "Authorization": "Bearer {token}", - "Content-Type": "application/json", - }}, - body: JSON.stringify({{ - model_path: "{GGUF_REPO}", - gguf_variant: "{GGUF_VARIANT}", - is_lora: false, - max_seq_length: 2048, - }}), - }}); - return {{status: r.status, body: await r.json()}}; - }}""") + # AbortSignal-bounded: the macos-14 --single-process Chromium had been + # observed wedging on this exact in-page fetch (run 25696797934 / job + # 75446949358) with zero further requests reaching the server. The + # 3-min budget is generous for a cold-cache GGUF load; on a wedge we + # surface a clean failure instead of a 30-min runner cancel. + load_resp = evaluate_fetch( + page, + f"{BASE}/api/inference/load", + method = "POST", + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + body = { + "model_path": GGUF_REPO, + "gguf_variant": GGUF_VARIANT, + "is_lora": False, + "max_seq_length": 2048, + }, + timeout_ms = LOAD_FETCH_TIMEOUT_MS, + ) + if load_resp.get("error"): + fail(f"/api/inference/load wedged: {load_resp['error']!r}") if load_resp["status"] != 200: fail( - f"/api/inference/load returned {load_resp['status']}: {load_resp.get('body')!r}" + f"/api/inference/load returned {load_resp['status']}: " + f"{load_resp.get('body')!r}" ) - info(f"loaded model: {load_resp['body'].get('display_name')}") + info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}") # Studio caches the per-context model state in zustand; reload # to make the chat composer pick up the loaded model. @@ -1177,10 +1214,13 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── # 14. /api/health stays healthy throughout. # ───────────────────────────────────────────────────── - health = page.evaluate(f"""async () => {{ - const r = await fetch("{BASE}/api/health"); - return {{status: r.status, body: await r.text()}}; - }}""") + health = evaluate_fetch( + page, + f"{BASE}/api/health", + timeout_ms = FETCH_TIMEOUT_MS, + ) + if health.get("error"): + fail(f"/api/health wedged: {health['error']!r}") if health["status"] != 200: fail(f"/api/health returned {health['status']}") @@ -1267,13 +1307,14 @@ with sync_playwright() as p: # The browser still has the pre-rotation access token. Refresh # tokens were revoked server-side by /change-password (auth.py), # so /api/auth/refresh from the browser context must now fail. - refresh_after = page.evaluate(f"""async () => {{ - const r = await fetch("{BASE}/api/auth/refresh", {{ - method: "POST", - credentials: "include", - }}); - return {{status: r.status}}; - }}""") + refresh_after = evaluate_fetch( + page, + f"{BASE}/api/auth/refresh", + method = "POST", + timeout_ms = FETCH_TIMEOUT_MS, + ) + if refresh_after.get("error"): + fail(f"/api/auth/refresh wedged: {refresh_after['error']!r}") if refresh_after["status"] == 200: fail(f"/api/auth/refresh should fail after CLI rotation; got 200") info( @@ -1384,4 +1425,5 @@ with sync_playwright() as p: ) info("PASS comprehensive UI flow") + _watchdog.cancel() browser.close() diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 92025ed555..dccd2e423d 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -40,7 +40,9 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) from _playwright_robust import ( # noqa: E402 chromium_launch_args, click_and_wait_for_response, + evaluate_fetch, install_view_transition_killer, + install_wall_clock_watchdog, is_benign_page_error, recover_or_replace_page, wait_for_health, @@ -59,6 +61,9 @@ STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1" # turn timeout because gemma-3-270m CPU inference is 3-5x slower than # ubuntu-latest's. TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000")) +WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720")) +FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000")) +LOAD_FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_LOAD_TIMEOUT_MS", "180000")) _n = [0] _failed: list[str] = [] @@ -94,6 +99,11 @@ def runtime_warn(m: str) -> None: with sync_playwright() as p: + _watchdog = install_wall_clock_watchdog( + WALL_TIMEOUT_S, + label = "ui-extra", + info = info, + ) # Health pre-flight (best-effort). Same rationale as in # playwright_chat_ui.py: bash-side health wait can succeed before # the auth DB has finished migrating on macos-14 free runners. @@ -261,36 +271,44 @@ with sync_playwright() as p: if not token: fail("no access token after change-password") sys.exit(1) - load_resp = page.evaluate(f"""async () => {{ - const r = await fetch("{BASE}/api/inference/load", {{ - method: "POST", - headers: {{ - "Authorization": "Bearer {token}", - "Content-Type": "application/json", - }}, - body: JSON.stringify({{ - model_path: "{GGUF_REPO}", - gguf_variant: "{GGUF_VARIANT}", - is_lora: false, - max_seq_length: 2048, - }}), - }}); - return {{status: r.status, body: await r.json()}}; - }}""") + load_resp = evaluate_fetch( + page, + f"{BASE}/api/inference/load", + method = "POST", + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + body = { + "model_path": GGUF_REPO, + "gguf_variant": GGUF_VARIANT, + "is_lora": False, + "max_seq_length": 2048, + }, + timeout_ms = LOAD_FETCH_TIMEOUT_MS, + ) + if load_resp.get("error"): + fail(f"/api/inference/load wedged: {load_resp['error']!r}") + sys.exit(1) if load_resp["status"] != 200: fail(f"/api/inference/load -> {load_resp['status']}: {load_resp.get('body')!r}") sys.exit(1) - info(f"loaded model: {load_resp['body'].get('display_name')}") + info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}") page.reload() composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) # Detect chat-only mode: /api/health.chat_only is the source of truth. # In chat-only mode, /studio + /export redirect to /chat. - health = page.evaluate(f"""async () => {{ - const r = await fetch("{BASE}/api/health"); - return await r.json(); - }}""") + health_resp = evaluate_fetch( + page, + f"{BASE}/api/health", + timeout_ms = FETCH_TIMEOUT_MS, + ) + if health_resp.get("error"): + fail(f"/api/health wedged: {health_resp['error']!r}") + sys.exit(1) + health = health_resp.get("body") or {} chat_only = bool(health.get("chat_only")) info(f"chat_only mode: {chat_only}") @@ -588,4 +606,5 @@ with sync_playwright() as p: info(f" - {m}") sys.exit(1) info("PASS extra UI flow") + _watchdog.cancel() browser.close()