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

Merge remote-tracking branch 'origin/main' into r5748

Browse source
This commit is contained in:
Daniel Han 2026-07-16 04:54:39 +00:00
commit 1c39a283f6
426 changed files with 66427 additions and 6621 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -376,7 +376,7 @@ case "$MODE" in
hermes) patch_hermes_tools none
invoke_via_connect "$OUT" -z "$PROMPT" ;;
openclaw) patch_openclaw_agent notools
invoke_via_connect "$OUT" agent --local --agent ci \
CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \
--model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;;
*) invoke_via_connect "$OUT" "$PROMPT" ;;
esac
@ -449,7 +449,7 @@ case "$MODE" in
fi ;;
opencode) invoke_via_connect "$out" run "$prompt" ;;
hermes) invoke_via_connect "$out" -z "$prompt" ;;
openclaw) invoke_via_connect "$out" agent --local --agent ci \
openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \
--model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;;
*) invoke_via_connect "$out" "$prompt" ;;
esac
@ -527,6 +527,154 @@ case "$MODE" in
echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
;;
# ── resume: does a launched agent's session survive exit and resume? ────
# Unlike the other modes, this drives the real LAUNCH path (`unsloth start
# <agent> ...`, the interactive default), not the --no-launch recipe. That
# path relocates each agent's home to a throwaway temp dir wiped on exit, so
# a session cannot be resumed -- unless --persist routes it to the stable
# Unsloth agents dir instead. We run one headless turn per pass and check
# whether the turn left a session in a persistent store (deterministic, no
# reliance on the model recalling anything), for a baseline pass and a
# --persist pass, and assert the expected split for this agent.
resume)
CODEWORD="PLATYPUS7"
T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK."
T2="What codeword did I ask you to remember? Reply with just that word."
WORK="$WORKDIR_BASE/${AGENT}-resume"
# STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to.
# Read it from a --no-launch probe (which also writes the agent's config
# there). codex/pi relocate their whole home/HOME here; opencode/claude keep
# their session data in a fixed user dir, so STABLE_HOME stays empty for them.
parse_connect
case "$AGENT" in
codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;;
pi) STABLE_HOME="$(raw_env HOME)" ;;
*) STABLE_HOME="" ;;
esac
# The persistent stores a session would land in if it were NOT wiped. We
# count files here before/after each turn; a positive delta means the
# session persisted (is resumable), zero means it went to a wiped temp dir.
resume_tracked_dirs() {
case "$AGENT" in
codex) printf '%s\n' "$HOME/.codex" ;;
opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;;
claude) printf '%s\n' "$HOME/.claude" ;;
pi) printf '%s\n' "$HOME/.pi" ;;
*) : ;;
esac
[ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME"
}
count_session_files() {
local total=0 d n
while IFS= read -r d; do
[ -n "$d" ] && [ -d "$d" ] || continue
n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n))
done < <(resume_tracked_dirs)
echo "$total"
}
# The headless first-turn subcommand per agent (mirrors file-edit's map),
# forwarded verbatim through the launch path as passthrough args.
set_t1_cmd() {
case "$AGENT" in
claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;;
codex) T1_CMD=(exec "$T1") ;;
opencode) T1_CMD=(run "$T1") ;;
pi) T1_CMD=(-p "$T1") ;;
*) guide_fail "resume mode does not cover agent '$AGENT'" ;;
esac
}
# Run one headless turn through the launch path. $1=outfile, $2="" or
# "--persist", rest = the agent subcommand. --yolo auto-approves so no tool
# prompt can hang; --api-key attaches to the already-served CI model.
launch_turn() {
local out="$1" rflag="$2"; shift 2
local flag=(); [ -n "$rflag" ] && flag=("$rflag")
run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \
--api-key "$UNSLOTH_API_KEY" "$@"
local rc=$?
redact "$out"
return "$rc"
}
# One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED
# from the session-store delta. Runs in the main shell (not a command
# substitution) so a hang's guide_fail actually fails the job and the
# progress lines reach the CI log. $1 = "" (baseline) or "--persist".
RESULT=""
run_pass() {
local rflag="$1" label="baseline"
[ -n "$rflag" ] && label="resume"
rm -rf "$WORK"; mkdir -p "$WORK"
set_t1_cmd
local out="$LOGS_DIR/${AGENT}-resume-${label}.txt"
local before after rc
before="$(count_session_files)"
pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK"
launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$?
popd >/dev/null || true
after="$(count_session_files)"
echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})"
# The turn must succeed for the delta to mean anything: an agent that writes a
# session file then errors would otherwise be misread as PERSISTED. Mirror the
# file-edit mode and fail the pass on a non-zero launch (the flagship codex recall
# below stays WARN-only, driven by its own launch_turn calls).
[ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \
guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; }
if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi
}
run_pass ""; BASELINE="$RESULT"
# Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix.
# opencode/claude persist either way, so the baseline already proves it and a
# second full CPU turn only risks a timeout; skip it for them.
case "$AGENT" in
codex|pi) run_pass "--persist"; RESUME="$RESULT" ;;
*) RESUME="n/a (persists either way)" ;;
esac
# Expected: codex/pi relocate their whole home to the temp dir, so a plain
# launch is WIPED and only --persist PERSISTS. opencode/claude keep their
# session data in a fixed user dir, so the baseline already PERSISTS.
case "$AGENT" in
codex|pi) EXPECT_BASELINE="WIPED" ;;
opencode|claude) EXPECT_BASELINE="PERSISTED" ;;
esac
echo "──────────────────────────────────────────────"
echo "[$AGENT] RESUME EXPERIMENT"
echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})"
echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}"
echo "──────────────────────────────────────────────"
[ "$BASELINE" = "$EXPECT_BASELINE" ] \
|| guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}"
case "$AGENT" in
codex|pi)
[ "$RESUME" = "PERSISTED" ] \
|| guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;;
esac
# Flagship behavioral proof (codex only, WARN-only): after a --persist plant,
# resume the session and check the model actually recalls the codeword. A
# miss is not a failure (the CI model is small); the mechanism gate above is
# the real assertion.
if [ "$AGENT" = "codex" ]; then
rm -rf "$WORK"; mkdir -p "$WORK"
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true
if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then
echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}"
else
echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed"
fi
fi
echo "[$AGENT] resume OK"
;;
*)
echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2
exit 2

4
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -272,6 +272,7 @@ jobs:
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
@ -361,10 +362,13 @@ jobs:
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \
tests/test_bad_mappings_redirect.py \
tests/test_prefetch_snapshot_scope.py \
tests/test_gemma_2b_mapper_key.py \
--deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
# The deselected test monkeypatches flash_attn_varlen_func, which is
# only bound on the module when `flash_attn` is importable. flash_attn

14
.github/workflows/cross-platform-parity-ci.yml vendored
View file

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Runs tests/python/test_cross_platform_parity.py on Windows and macOS.
# Runs installer parity and autostart opt-out tests on Windows and macOS.
#
# Why: that test is the guard that install.sh and install.ps1 stay in
# sync, but today it only runs on ubuntu-latest (auto-discovered by
@ -21,6 +21,7 @@ on:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
@ -28,6 +29,7 @@ on:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
@ -57,5 +59,11 @@ jobs:
python-version: '3.12'
cache: 'pip'
- run: python -m pip install -U pip pytest
- name: Cross-platform parity test
run: python -m pytest tests/python/test_cross_platform_parity.py -q
- name: Cross-platform parity tests
env:
UNSLOTH_NO_TORCH: '1'
run: >-
python -m pytest
tests/python/test_cross_platform_parity.py
tests/test_installer_skip_autostart.py
-q

170
.github/workflows/local-agent-guides-ci.yml vendored
View file

@ -471,6 +471,176 @@ jobs:
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job: resume
# Does a conversation started with `unsloth start <agent>` survive exit
# and resume? This drives the REAL launch path (not the --no-launch
# recipe the other jobs use). A plain launch relocates the agent home to
# a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the
# session to the stable Unsloth agents dir so it persists. opencode/claude
# keep their session data in a fixed user dir, so they persist either way.
# Dispatch-only: it is an end-to-end experiment, not a PR gate.
# ═════════════════════════════════════════════════════════════════════
resume:
name: resume (${{ matrix.agent }})
if: github.event_name == 'workflow_dispatch'
runs-on: ubuntu-latest
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
# codex/pi relocate their whole home (resume broken without --persist);
# opencode/claude keep session data in a fixed dir (resume already works).
# One agent from each class proves the split end to end; openclaw/hermes
# share codex's relocation mechanism and are covered by the unit tests.
agent: [codex, opencode, claude, pi]
env:
GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF
GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf
STUDIO_PORT: '18904'
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- name: Linux deps for llama.cpp prebuilt
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libcurl4-openssl-dev libssl-dev jq
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- name: Restore GGUF model file
id: cache-gguf
uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
continue-on-error: true
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Download GGUF if cache miss
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
- name: Save GGUF model file
if: always() && steps.download-gguf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
unsloth studio reset-password
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
--extra "--seed $UNSLOTH_SEED --temp 0" \
--health-timeout 900
- name: Preflight the agent's API dialect (class-a isolation)
env:
AGENT: ${{ matrix.agent }}
run: |
set -uo pipefail
B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY"
preflight_fail() {
echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**.";
exit 1
}
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \
-H "Authorization: Bearer $K") || true
[ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code"
case "$AGENT" in
claude)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code"
;;
codex)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true
[ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code"
;;
*)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code"
;;
esac
echo "preflight OK for $AGENT"
- name: Install agent CLI (class-b isolation)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-install.sh "$AGENT"
- name: Resume experiment (launch path)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT"
- name: Collect server logs (debug)
if: always()
run: |
mkdir -p logs/studio-logs
cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
done
fi
- name: Stop Studio
if: always()
run: |
if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then
kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true
fi
sleep 2
ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true
- name: Upload logs
if: always()
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: resume-${{ matrix.agent }}-log
path: |
logs/
agent-workdir/
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job 3: prompt-cache
# (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0

6
.github/workflows/security-audit.yml vendored
View file

@ -2,8 +2,8 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Multi-language supply-chain audit. Triggers:
# - PRs touching any dependency manifest (Python / npm / Cargo) or
# this workflow file,
# - PRs touching any dependency manifest (Python / npm / Cargo), a
# scanner or its allowlist baseline, or this workflow file,
# - push to main / pip,
# - nightly @ 04:13 UTC so newly-published advisories surface even
# when no PR opens,
@ -57,7 +57,9 @@ on:
- 'studio/src-tauri/Cargo.lock'
- 'pyproject.toml'
- 'scripts/scan_packages.py'
- 'scripts/scan_packages_baseline.json'
- 'scripts/scan_npm_packages.py'
- 'scripts/scan_npm_packages_baseline.json'
- '.github/workflows/security-audit.yml'
push:
branches: [main, pip]

181
.github/workflows/studio-inference-smoke.yml vendored
View file

@ -444,6 +444,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -464,10 +466,26 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
def post_sse(path, body, *, timeout = 600, retries = 1, complete_on = None):
"""POST a streaming request and accumulate the assistant
text deltas. The server-side agentic loop ALWAYS returns
SSE regardless of the request's `stream` field, so any
@ -483,6 +501,22 @@ jobs:
invocation markers / tool output, since
`delta.content` alone is not evidence
that the tool path executed.
A shared CI runner can stall the stream transport (the
connection opening, or a mid-stream read) even when Studio
is healthy, so retry a stall once with a fresh request
capped at 300s. A stall means the stream did NOT complete,
so partial events are normally NOT returned (an early
tool_start with no tool_end is not proof the tool loop
finished). The one exception is `complete_on`: an optional
predicate over the events collected so far -- when a stall
happens after it is already satisfied (the tool ran and
produced its result before the trailing read timed out),
those events are returned rather than discarded, so the
stall-after-answer case still counts. HTTP status errors
surface immediately; a stall that yields no completed result
across all attempts re-raises so the caller can rotate to
the next seed.
"""
body = {**body, "stream": True}
data = json.dumps(body).encode()
@ -495,26 +529,45 @@ jobs:
"Content-Type": "application/json",
},
)
parts = []
events = []
with urllib.request.urlopen(req, timeout = timeout) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
events.append(payload)
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts), events
for attempt in range(retries + 1):
parts = []
events = []
t = timeout if attempt == 0 else min(timeout, 300)
try:
with urllib.request.urlopen(req, timeout = t) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
events.append(payload)
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts), events
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# A stall after the tool already produced its result is
# the case this probe exists to tolerate: keep those
# events. But a stall with only an early tool_start (no
# completed output) is not proof the tool loop finished,
# so it must not pass -- retry once, then raise so
# _run_tool_probe rotates to the next seed.
if complete_on is not None and complete_on(events):
print(f"[retry-sse] {path}: {exc!r}; keeping {len(events)} completed events", flush = True)
return "".join(parts), events
if attempt == retries:
raise
print(f"[retry-sse] {path}: {exc!r}", flush = True)
time.sleep(15)
_STUDIO_TOOL_TYPES = {
"tool_start", "tool_end", "tool_use", "tool_result",
@ -651,17 +704,54 @@ jobs:
"""
attempts_log = []
best = None
# Cap the wall-clock spent rotating through stalled seeds so a
# persistent no-data wedge fails fast (clean assertion) instead
# of being killed by the job's timeout-minutes. A healthy or
# merely degenerate round answers in seconds, so all seeds still
# run in the normal case; only stalls consume the budget.
probe_deadline = time.monotonic() + 300
for attempt_i in range(max_attempts):
# Cap each read by the budget still remaining (not just a flat
# 180s) and skip an attempt too small to finish, so the whole
# rotation stays within ~300s -- two probes then fit the job's
# timeout-minutes even if every seed stalls.
remaining = int(probe_deadline - time.monotonic())
if attempt_i and remaining < 30:
print(f"[tools] {label}: seed-rotation budget spent after {attempt_i} attempts", flush = True)
break
attempt_seed = SEED + attempt_i
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": prompt}],
"enable_tools": True,
"enabled_tools": enabled,
"session_id": f"{session}-att{attempt_i}",
"temperature": TOOL_PROBE_TEMP,
"seed": attempt_seed,
"max_tokens": 600,
})
try:
# Bounded per-attempt timeout, no inner retry -- the seed
# loop IS the retry, so a stall raises quickly and rotates
# rather than spending post_sse's full 600+300s. complete_on
# keeps a stall that already produced the tool result (only
# the trailing read timed out) instead of discarding it.
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": prompt}],
"enable_tools": True,
"enabled_tools": enabled,
"session_id": f"{session}-att{attempt_i}",
"temperature": TOOL_PROBE_TEMP,
"seed": attempt_seed,
"max_tokens": 600,
}, timeout = min(180, remaining), retries = 0,
complete_on = lambda ev: _tool_invoked(ev) and _tool_output_contains(ev, *needles))
except urllib.error.HTTPError:
# HTTPError subclasses URLError, so re-raise a real 4xx/5xx
# here instead of letting the transport-stall handler below
# swallow it and rotate seeds -- an endpoint status failure
# must surface, not be masked as missing tool evidence.
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# A transport stall that outlived post_sse's own retry:
# log it as a failed attempt and rotate to the next seed
# rather than sinking the whole probe on one bad stream.
attempts_log.append({
"attempt": attempt_i, "seed": attempt_seed,
"transport_error": repr(exc),
})
print(f"[tools] retry {label} attempt {attempt_i}: transport {exc!r}", flush = True)
continue
invoked = _tool_invoked(events)
produced = _tool_output_contains(events, *needles)
attempts_log.append({
@ -722,6 +812,9 @@ jobs:
# enough that requiring a tool_call marker would create
# red-herring failures from infra rather than from Studio.
try:
# Best-effort and bounded: a single 180s attempt keeps a stall
# from eating the job's timeout-minutes (it already WARNs, so a
# retry buys nothing).
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
@ -730,7 +823,7 @@ jobs:
"temperature": 0.0,
"seed": SEED,
"max_tokens": 400,
})
}, timeout = 180, retries = 0)
print(
f"[tools] PASS web_search stream ({len(content)} chars in content, "
f"{len(events)} raw events)"
@ -938,6 +1031,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -956,8 +1051,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
# llama.cpp's HTTP server supports OpenAI-compatible JSON

126
.github/workflows/studio-mac-inference-smoke.yml vendored
View file

@ -430,6 +430,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -450,14 +452,41 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
def post_sse(path, body, *, timeout = 600, retries = 1, soft = False):
"""POST a streaming request and accumulate the assistant
text deltas. The server-side agentic loop ALWAYS returns
SSE regardless of the request's `stream` field, so any
call with enable_tools=true must use this helper."""
call with enable_tools=true must use this helper.
A shared CI runner can stall the stream transport (the
connection opening, or a mid-stream read) even when Studio
is healthy, so harden the read three ways: retry a stall
once with a fresh request capped at 300s; return any text
already streamed before a stall (a stall on the trailing
tokens, after the answer arrived, still counts); and when
every attempt yields nothing, a hard call re-raises while a
soft call (the best-effort server-side tool probes) returns
None so the caller can WARN instead of sinking the whole
job. HTTP status errors always surface immediately."""
body = {**body, "stream": True}
data = json.dumps(body).encode()
req = urllib.request.Request(
@ -469,24 +498,43 @@ jobs:
"Content-Type": "application/json",
},
)
parts = []
with urllib.request.urlopen(req, timeout = timeout) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
for attempt in range(retries + 1):
parts = []
t = timeout if attempt == 0 else min(timeout, 300)
try:
with urllib.request.urlopen(req, timeout = t) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# Text already streamed is a valid signal -- keep it
# rather than re-running a heavy generation.
if parts:
joined = "".join(parts)
print(f"[retry-sse] {path}: {exc!r}; keeping {len(joined)} partial chars", flush = True)
return joined
if attempt == retries:
if soft:
print(f"[tools] WARN {path}: SSE transport stalled with no data ({exc!r}) -- non-blocking", flush = True)
return None
raise
print(f"[retry-sse] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. Standard OpenAI function calling ──────────────────────
weather_tool = {
@ -557,6 +605,10 @@ jobs:
# macos-14 free runner is ~10 tok/s on Qwen3.5-2B Q4_K_XL;
# cap max_tokens tightly so each SSE round stays under ~30s
# even when the model stalls in a degenerate output state.
# retries=0 on the best-effort probes: this job's 25-minute cap
# allows a 10-minute model load, so a no-data stall must be a
# single 180s attempt (not 180+15+180s) to leave room for the
# thinking checks. A soft/best-effort probe only WARNs anyway.
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}],
"enable_tools": True,
@ -565,8 +617,10 @@ jobs:
"temperature": TEMP,
"seed": SEED,
"max_tokens": 128,
}, timeout = 180)
if "56088" in content or "56,088" in content:
}, timeout = 180, retries = 0, soft = True)
if content is None:
print("[tools] WARN python tool: SSE transport stalled after retries -- non-blocking")
elif "56088" in content or "56,088" in content:
print(f"[tools] PASS python tool ({len(content)} chars, found 56088)")
else:
# Empty stream is a known Mac-quant degeneracy too; log
@ -598,7 +652,7 @@ jobs:
"temperature": TEMP,
"seed": SEED,
"max_tokens": 96,
}, timeout = 180)
}, timeout = 180, retries = 0)
print(f"[tools] PASS web_search stream ({len(content)} chars)")
except Exception as exc:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
@ -825,6 +879,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -848,8 +904,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
# llama.cpp's HTTP server supports OpenAI-compatible JSON

296
.github/workflows/studio-windows-inference-smoke.yml vendored
View file

@ -634,6 +634,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -656,10 +658,41 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
def post_sse(path, body, *, timeout = 600, retries = 1, soft = False):
# The server-side agentic loop always answers over SSE. A
# shared CI runner can stall the stream transport (the
# connection opening, or a mid-stream read) even when Studio
# is healthy, so harden the read three ways:
# * retry a transport stall once with a fresh request,
# capped at 300s (a healthy server answers a retry
# quickly, a wedged one never does);
# * return any text already streamed before a stall, so a
# stall on the trailing tokens -- after the answer
# arrived -- still counts;
# * when every attempt yields nothing, a hard call
# re-raises while a soft call (the best-effort
# server-side tool probes) returns None so the caller
# can WARN instead of sinking the whole job.
# HTTP status errors always surface immediately.
body = {**body, "stream": True}
data = json.dumps(body).encode()
req = urllib.request.Request(
@ -671,24 +704,43 @@ jobs:
"Content-Type": "application/json",
},
)
parts = []
with urllib.request.urlopen(req, timeout = timeout) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
for attempt in range(retries + 1):
parts = []
t = timeout if attempt == 0 else min(timeout, 300)
try:
with urllib.request.urlopen(req, timeout = t) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# Text already streamed is a valid signal -- keep it
# rather than re-running a heavy generation.
if parts:
joined = "".join(parts)
print(f"[retry-sse] {path}: {exc!r}; keeping {len(joined)} partial chars", flush = True)
return joined
if attempt == retries:
if soft:
print(f"[tools] WARN {path}: SSE transport stalled with no data ({exc!r}) -- non-blocking", flush = True)
return None
raise
print(f"[retry-sse] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. Standard OpenAI function calling ──────────────────────
weather_tool = {
@ -731,6 +783,11 @@ jobs:
)
# ── 2. Server-side python tool ───────────────────────────────
# Bound each soft probe to a single 180s attempt (timeout=180,
# retries=0): this job runs two of them back-to-back under a
# 30-minute cap, so the default 600+15+300s per stall could hit
# the workflow timeout before the thinking checks run. A soft
# probe only WARNs anyway, so a retry buys nothing.
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}],
"enable_tools": True,
@ -739,8 +796,10 @@ jobs:
"temperature": TEMP,
"seed": SEED,
"max_tokens": 600,
})
if "56088" in content or "56,088" in content:
}, timeout = 180, retries = 0, soft = True)
if content is None:
print("[tools] WARN python tool: SSE transport stalled after retries -- non-blocking")
elif "56088" in content or "56,088" in content:
print(f"[tools] PASS python tool ({len(content)} chars, found 56088)")
else:
assert content, "python tool: SSE stream empty"
@ -762,8 +821,10 @@ jobs:
"temperature": TEMP,
"seed": SEED,
"max_tokens": 600,
})
if "hello-bash-tool" in content:
}, timeout = 180, retries = 0, soft = True)
if content is None:
print("[tools] WARN terminal tool: SSE transport stalled after retries -- non-blocking")
elif "hello-bash-tool" in content:
print(f"[tools] PASS terminal tool ({len(content)} chars)")
else:
assert content, "terminal tool: SSE stream empty"
@ -784,7 +845,7 @@ jobs:
"temperature": TEMP,
"seed": SEED,
"max_tokens": 400,
})
}, timeout = 180, retries = 0)
print(f"[tools] PASS web_search stream ({len(content)} chars)")
except Exception as exc:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
@ -1063,6 +1124,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -1082,8 +1145,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
status, data = post("/v1/chat/completions", {
@ -1334,42 +1413,75 @@ jobs:
try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { }
}
- name: Hide Visual Studio + CMake (simulate a host with no build tools)
- name: Prepare no-build-tools simulation
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# A Program Files dir can hold a transient handle (Defender / MSBuild node)
# so Rename-Item intermittently fails with "Access is denied"; retry to ride it out.
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
$pf = Join-Path $root 'ProgramFiles'
$pfx86 = Join-Path $root 'ProgramFilesx86'
New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null
$blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
foreach ($tool in @('cmake', 'cl.exe')) {
foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) {
if ($cmd.Source) {
$dir = Split-Path -Parent $cmd.Source
if ($dir) {
[void] $blocked.Add(
[Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\'))
}
}
}
}
# Rename the Visual Studio install roots (incl. the Installer that holds
# vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss.
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) {
Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff')
Write-Host "Hid VS: $d"
}
# Normalized comparison so registry spellings (trailing slash,
# unexpanded %VAR%) still match.
function Test-Blocked([string]$p) {
$n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\')
return $blocked.Contains($n)
}
# Surgically rename each cmake executable on PATH (not its parent dir --
# cmake can share a dir with other shims) so Get-Command cmake fails.
$hidden = @()
foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) {
if ($c.Source -and (Test-Path -LiteralPath $c.Source)) {
Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off')
$hidden += $c.Source
Write-Host "Hid cmake: $($c.Source)"
}
$pathParts = $env:Path -split [IO.Path]::PathSeparator |
Where-Object { $_ -and -not (Test-Blocked $_) }
$noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator
# install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment
# rebuild the session Path from these scopes mid-install, so filter
# them too. Originals are saved for the cleanup step.
foreach ($scope in @('Machine', 'User')) {
$orig = [Environment]::GetEnvironmentVariable('Path', $scope)
if (-not $orig) { continue }
Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline
$kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';'
[Environment]::SetEnvironmentVariable('Path', $kept, $scope)
Write-Host "Filtered $scope Path scope."
}
"NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
Write-Host "ProgramFiles simulation root: $pf"
Write-Host "ProgramFiles(x86) simulation root: $pfx86"
if ($blocked.Count -gt 0) {
Write-Host "Removed build-tool PATH dirs:"
$blocked | Sort-Object | ForEach-Object { Write-Host " $_" }
} else {
Write-Host "No cmake or cl.exe PATH dirs found to remove."
}
("HIDDEN_CMAKE=" + ($hidden -join '|')) | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Assert Visual Studio + CMake are genuinely undetectable
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Set in-script: the runner does not apply step-level env keys with
# parentheses (`ProgramFiles(x86)`), so vswhere still found VS.
if (-not $env:NO_BUILD_TOOLS_PROGRAMFILES) { Write-Error "NO_BUILD_TOOLS_* env missing (Prepare step did not run?)"; exit 1 }
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
. (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1')
$setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1'
foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Find-VsBuildTools')) {
@ -1394,6 +1506,10 @@ jobs:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
# Set in-script (see the assert step); child processes inherit these.
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
New-Item -ItemType Directory -Force -Path logs | Out-Null
$ProgressPreference = 'SilentlyContinue'
& ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log
@ -1480,19 +1596,19 @@ jobs:
[ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; }
echo "Inference OK without Visual Studio: $CONTENT"
- name: Restore Visual Studio + CMake
- name: Clean no-build-tools simulation
if: always()
shell: pwsh
run: |
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
$off = "$d.vsoff"
if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" }
}
if ($env:HIDDEN_CMAKE) {
foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) {
if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
foreach ($scope in @('Machine', 'User')) {
$saved = Join-Path $root "orig-path-$scope.txt"
if (Test-Path -LiteralPath $saved) {
[Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope)
Write-Host "Restored $scope Path scope."
}
}
Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue
- name: Stop Studio
if: always()
@ -1540,21 +1656,34 @@ jobs:
with:
python-version: '3.12'
- name: Hide Visual Studio
- name: Prepare no-build-tools simulation
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Retry the rename: a Program Files dir can hold a transient handle that
# makes Rename-Item intermittently fail with "Access is denied".
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
$pf = Join-Path $root 'ProgramFiles'
$pfx86 = Join-Path $root 'ProgramFilesx86'
New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null
$blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
foreach ($tool in @('cmake', 'cl.exe')) {
foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) {
if ($cmd.Source) {
$dir = Split-Path -Parent $cmd.Source
if ($dir) { [void] $blocked.Add($dir) }
}
}
}
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
}
$pathParts = $env:Path -split [IO.Path]::PathSeparator |
Where-Object { $_ -and -not $blocked.Contains($_) }
$noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator
"NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS)
env:
@ -1577,25 +1706,34 @@ jobs:
echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling."
- name: The prebuilt resolver runs without Visual Studio
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
$ErrorActionPreference = 'Stop'
# pwsh: bash cannot export `ProgramFiles(x86)`; set in-script so the
# python child inherits the overrides.
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
# Resolver-only (no GPU on hosted runners, so the host resolves to the
# CPU bundle). The point is that resolution needs no compiler/VS.
python -m pip install --upgrade huggingface_hub
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > /tmp/resolve.json || {
echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; }
cat /tmp/resolve.json
echo "Prebuilt resolver ran with no Visual Studio present."
if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 }
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json
if ($LASTEXITCODE -ne 0) {
Write-Host "::error::resolver exited non-zero"
if (Test-Path resolve.json) { Get-Content resolve.json }
exit 1
}
Get-Content resolve.json
Write-Host "Prebuilt resolver ran with no Visual Studio present."
- name: Restore Visual Studio
- name: Clean no-build-tools simulation
if: always()
shell: pwsh
run: |
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
$off = "$d.vsoff"
if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" }
}
Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue
# ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ──
pester:

6
.github/workflows/studio-windows-update-smoke.yml vendored
View file

@ -6,9 +6,9 @@
# windows-latest runner:
#
# 1. install.ps1 --local --no-torch installs Studio AND auto-fetches
# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu-
# x64 from ggml-org/llama.cpp). Hitting the source-build fallback
# is treated as an Unsloth bug -- Studio must always pick the
# the prebuilt llama.cpp Windows binary (app-<tag>-windows-x64-cpu
# from unslothai/llama.cpp). Hitting the source-build fallback is
# treated as an Unsloth bug -- Studio must always pick the
# prebuilt on Windows.
# 2. unsloth studio update --local is idempotent. Two consecutive
# runs both report "prebuilt up to date and validated", no

86
.github/workflows/version-compat-ci.yml vendored
View file

@ -285,6 +285,92 @@ jobs:
tests/vllm_compat/test_extended_module_imports.py \
-v --tb=short
# Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike
# the static symbol/source greps above, this drives unsloth's actual
# source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only
# runner under the tests/conftest.py spoof harness -- no GPU, no training.
# Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple
# per-token-logps return, restructured PEFT ref-adapter block) by asserting
# the generated Unsloth trainer still satisfies the transform contracts.
grpo-fake-run:
name: GRPO fake-run (latest + main TRL, CPU spoof)
runs-on: ubuntu-latest
timeout-minutes: 18
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
path: unsloth
- name: Clone unsloth-zoo @ main
run: |
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'
cache: 'pip'
- name: Install CPU torch + ecosystem + TRL latest
run: |
python -m pip install --upgrade pip
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10'
# Ecosystem floors unsloth needs; TRL itself is installed last so it
# can pull the transformers/peft it requires.
pip install \
'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \
'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \
'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow
pip install --upgrade trl
pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo"
pip install --no-deps -e ./unsloth
- name: Fake-run vs TRL latest
env:
UNSLOTH_IS_PRESENT: '1'
UNSLOTH_COMPILE_DISABLE: '1'
# Disable dynamo/inductor at the process level, before conftest.py's early
# `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner
# (defense in depth; the CPU fake-train also flips this at runtime).
TORCHDYNAMO_DISABLE: '1'
TORCH_COMPILE_DISABLE: '1'
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
run: |
cd unsloth
python -c "import trl; print('Resolved TRL', trl.__version__)"
PYTHONPATH=. python -m pytest \
tests/version_compat/test_trl_grpo_fake_run.py \
tests/version_compat/test_trl_fake_train_cpu.py \
-v --tb=short
# `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge
# TRL break does not red every PR. github.event_name is valid in a step if.
- name: Fake-run vs TRL main (scheduled / dispatch only)
if: ${{ github.event_name != 'pull_request' }}
env:
UNSLOTH_IS_PRESENT: '1'
UNSLOTH_COMPILE_DISABLE: '1'
TORCHDYNAMO_DISABLE: '1'
TORCH_COMPILE_DISABLE: '1'
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
run: |
pip install --upgrade "git+https://github.com/huggingface/trl"
cd unsloth
python -c "import trl; print('Resolved TRL', trl.__version__)"
PYTHONPATH=. python -m pytest \
tests/version_compat/test_trl_grpo_fake_run.py \
tests/version_compat/test_trl_fake_train_cpu.py \
-v --tb=short
# Daily-only: same suites but with --strict on importable upstream
# tags. Schedule-only so PR jobs stay fast; cron tolerates a flake.
daily-fresh-fetch:

25
README.md
View file

@ -84,7 +84,7 @@ Use the same command to update.
```bash
unsloth studio -p 8888
```
For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally.
For LAN or cloud access, add `-H 0.0.0.0` (raw port only; add `--cloudflare` for a public URL). By default, Unsloth is accessible only locally.
To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below).
@ -212,10 +212,23 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i
```bash
unsloth studio --secure -p 8888
```
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network.
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network (subject to your firewall). It does not create a public internet URL; add `--cloudflare` to also publish an internet-reachable `https://*.trycloudflare.com` link even behind a firewall. Only use this on a network you trust.
```bash
unsloth studio -H 0.0.0.0 -p 8888
```
The Cloudflare tunnel is **off by default**: `-H 0.0.0.0` exposes the raw port only, not a public internet URL. Pair the wildcard bind with `--cloudflare` (`unsloth studio -H 0.0.0.0 --cloudflare`) to also publish a public `https://*.trycloudflare.com` link, or prefer `--secure` (above), which keeps the raw port private. `--cloudflare` has no effect on a loopback bind.
The first time Studio is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Studio shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`):
```bash
unsloth studio --secure --password 'your-strong-password' # visible in `ps`/history
UNSLOTH_STUDIO_PASSWORD='your-strong-password' unsloth studio --secure # via env var
printf '%s\n' 'your-strong-password' | unsloth studio --secure --password - # via stdin
```
A literal `--password VALUE` is visible in the process list and shell history, so prefer the `UNSLOTH_STUDIO_PASSWORD` env var or `--password -` (stdin) for automation. This applies to any launch (public or a headless `-H 0.0.0.0` bind), and the password is set in the parent before the server binds, so it never reaches a re-executed child process.
Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Studio.
@ -230,6 +243,14 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh
$env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex
```
Skip the post-install prompt that starts Studio (useful for automated installs):
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh
```
```powershell
$env:UNSLOTH_SKIP_AUTOSTART=1; irm https://unsloth.ai/install.ps1 | iex
```
Pin the Python version:
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh

51
install.ps1
View file

@ -6,6 +6,7 @@
# irm | iex cannot forward arguments, so web installs take options as env vars set
# before the pipe (flags still work via .\install.ps1):
# $env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex # skip PyTorch (GGUF-only)
# $env:UNSLOTH_SKIP_AUTOSTART=1; irm https://unsloth.ai/install.ps1 | iex # do not prompt to launch
# $env:UNSLOTH_PYTHON='3.12'; irm https://unsloth.ai/install.ps1 | iex # pin Python version
# $env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex
# .\install.ps1 --no-torch # equivalent flag
@ -90,6 +91,7 @@ function Install-UnslothStudio {
if ($TauriMode) {
exit $Code
}
throw $Message
}
# ── Parse flags ──
@ -98,6 +100,7 @@ function Install-UnslothStudio {
$RepoRoot = ""
$TauriMode = $false
$SkipTorch = $false
$SkipAutostart = $false
$ShortcutsOnly = $false
$WithLlamaCppDir = ""
$argList = $args
@ -130,6 +133,7 @@ function Install-UnslothStudio {
# Env-var equivalent for web installs; an explicit flag still wins.
if ($env:UNSLOTH_NO_TORCH -in @('1', 'true', 'yes', 'on')) { $SkipTorch = $true }
if ($env:UNSLOTH_SKIP_AUTOSTART -in @('1', 'true', 'yes', 'on')) { $SkipAutostart = $true }
# Propagate to child processes so they also respect verbose mode.
# Process-scoped -- does not persist.
@ -469,6 +473,17 @@ function Install-UnslothStudio {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when the command pins an index, clear every uv index env var so
# it wins, then restore in finally. Other installs keep the user's mirror.
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@ -488,6 +503,7 @@ function Install-UnslothStudio {
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
}
}
@ -2160,7 +2176,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2174,7 +2190,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2205,7 +2221,7 @@ exit 0
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
# still completes; Studio setup retries ROCm afterwards.
@ -2214,7 +2230,7 @@ exit 0
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2228,7 +2244,7 @@ exit 0
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2240,7 +2256,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2252,7 +2268,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2280,7 +2296,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@ -2311,7 +2327,7 @@ exit 0
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
# is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
if (-not $SkipTorch) {
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
@ -2327,7 +2343,7 @@ exit 0
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
@ -2336,7 +2352,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
@ -2605,9 +2621,10 @@ exit 0
# Diagnostic only; never block install on a probe failure.
}
# In interactive terminals, ask the user before starting Studio.
# In interactive terminals, ask the user before starting Studio unless the
# caller explicitly disabled the post-install prompt.
# In non-interactive environments (CI, Docker) just print instructions.
$IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
$IsInteractive = (-not $SkipAutostart) -and [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
if ($IsInteractive) {
Write-Host ""
$reply = Read-Host " Start Unsloth Studio now? [Y/n]"
@ -2616,8 +2633,8 @@ exit 0
} else {
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
Write-Host ""
}
} else {
@ -2637,8 +2654,8 @@ exit 0
substep "& $_actLiteral"
substep "unsloth studio -p 8888"
}
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
Write-Host ""
}
}

59
install.sh
View file

@ -8,8 +8,9 @@
#
# Piped installs take options as env vars after the pipe (a bare `| sh --no-torch`
# makes sh reject --no-torch as its own option). Flags still work via ./install.sh:
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh # skip PyTorch (GGUF-only)
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh # pin Python version
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh # skip PyTorch (GGUF-only)
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh # do not prompt to launch
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh # pin Python version
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh
# Equivalent flags: ./install.sh --no-torch --python 3.12 (or pipe them: sh -s -- --no-torch)
#
@ -49,6 +50,7 @@ PACKAGE_NAME="unsloth"
TAURI_MODE=false
_USER_PYTHON=""
_NO_TORCH_FLAG=false
_SKIP_AUTOSTART=false
_VERBOSE=false
_SHORTCUTS_ONLY=false
_next_is_package=false
@ -88,6 +90,7 @@ done
# Env-var equivalents for piped installs; an explicit flag still wins.
case "${UNSLOTH_NO_TORCH:-}" in 1|true|TRUE|yes|YES|on|ON) _NO_TORCH_FLAG=true ;; esac
case "${UNSLOTH_SKIP_AUTOSTART:-}" in 1|true|TRUE|yes|YES|on|ON) _SKIP_AUTOSTART=true ;; esac
[ -z "$_USER_PYTHON" ] && [ -n "${UNSLOTH_PYTHON:-}" ] && _USER_PYTHON="$UNSLOTH_PYTHON"
if [ "$_VERBOSE" = true ]; then
@ -159,6 +162,12 @@ run_maybe_quiet() {
run_install_cmd() {
_label="$1"
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
esac
if _is_verbose; then
"$@" && return 0
_rc=$?
@ -1625,6 +1634,7 @@ _maybe_reroute_strixhalo_to_2404() {
# Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the
# GPU instead of falling back to the desktop-app prompt path.
[ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1"
[ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1"
_rr_args=""
[ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")"
[ -n "$_USER_PYTHON" ] && _rr_args="$_rr_args --python $(_rr_q "$_USER_PYTHON")"
@ -2210,9 +2220,9 @@ _expected_torch_flavor_tag() {
esac
}
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX /
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
# resolves (torch + every transitive dep) via --default-index -- the same URLs the
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
@ -2726,7 +2736,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2741,7 +2751,7 @@ if [ "$_MIGRATED" = true ]; then
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-}
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2764,7 +2774,7 @@ if [ "$_MIGRATED" = true ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2890,7 +2900,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@ -2913,18 +2923,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@ -2945,7 +2955,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2963,7 +2973,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
--upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git ${_ZOO_REF}..."
@ -2984,7 +2994,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2995,7 +3005,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git ${_ZOO_REF}..."
@ -3019,14 +3029,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
# Repair when flavor is wrong AND the index is plain --default-index reinstallable
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
@ -3037,7 +3047,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi
@ -3237,9 +3247,10 @@ printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!"
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""
# In interactive terminals, ask the user before starting Studio.
# In interactive terminals, ask the user before starting Studio unless the
# caller explicitly disabled the post-install prompt.
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
if [ -t 1 ]; then
if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then
echo ""
printf " Start Unsloth Studio now? [Y/n] "
# No readable answer (closed/EOF tty) defaults to no; Enter is still yes.
@ -3275,8 +3286,8 @@ if [ -t 1 ]; then
*)
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
echo ""
;;
esac
@ -3297,7 +3308,7 @@ else
substep "source $_li_act_q"
substep "unsloth studio -p 8888"
fi
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
echo ""
fi

7
pyproject.toml
View file

@ -42,6 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"}
include-package-data = true
[tool.setuptools.package-data]
unsloth_cli = ["codex_fallback_prompt.md"]
studio = [
"*.sh",
"*.ps1",
@ -73,7 +74,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.7.1",
"unsloth_zoo>=2026.7.3",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -94,7 +95,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.7.1",
"unsloth_zoo>=2026.7.3",
"torchvision",
"unsloth[triton]",
]
@ -579,7 +580,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.7.1",
"unsloth_zoo>=2026.7.3",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

326
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

55
studio/backend/auth/storage.py
View file

@ -18,6 +18,10 @@ from utils.paths import auth_db_path, ensure_dir
DB_PATH = auth_db_path()
DEFAULT_ADMIN_USERNAME = "unsloth"
# Single source for the password policy; models/auth.py ChangePasswordRequest
# and the terminal prompt both enforce it. Keep the unsloth_cli mirror in sync.
MIN_PASSWORD_LENGTH = 8
# Plaintext bootstrap password file beside auth.db, deleted on first password
# change so the credential never lingers on disk.
_BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password"
@ -79,11 +83,42 @@ def _load_bootstrap_password() -> Optional[str]:
def clear_bootstrap_password() -> None:
"""Delete the persisted bootstrap password file (called after password change)."""
"""Delete the persisted bootstrap password file (after a password change).
Best-effort: the new hash is already committed, so a locked/undeletable file
(Windows AV, read-only auth dir) must not fail the change.
"""
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
_BOOTSTRAP_PW_PATH.unlink(missing_ok = True)
try:
_BOOTSTRAP_PW_PATH.unlink(missing_ok = True)
except OSError as e:
# Removal failed (Windows AV, read-only auth dir). The hash is already
# committed, so don't fail the change -- but truncate the file so its
# stale plaintext can't be re-seeded by generate_bootstrap_password()
# if a later reset-password deletes auth.db and re-validates it.
try:
_BOOTSTRAP_PW_PATH.write_text("")
cleared = True
except OSError:
cleared = False
import sys
if cleared:
message = (
f"Warning: could not delete {_BOOTSTRAP_PW_PATH.name} ({e}); "
"cleared its contents so the old bootstrap password cannot be reused."
)
else:
# Neither removed nor truncated: stale plaintext is still on disk
# and would be reused if auth.db is reset. Don't claim otherwise.
message = (
f"Warning: could not delete or clear {_BOOTSTRAP_PW_PATH.name} ({e}); "
"its old bootstrap password is still on disk. Remove it manually to "
"prevent reuse after a reset."
)
print(message, file = sys.stderr, flush = True)
def _hash_token(token: str) -> str:
@ -547,8 +582,18 @@ def ensure_default_admin() -> bool:
return False
def update_password(username: str, new_password: str) -> bool:
"""Update password, clear first-login requirement, rotate JWT secret."""
def update_password(
username: str,
new_password: str,
*,
revoke_refresh_tokens: bool = False,
) -> bool:
"""Update password, clear first-login requirement, rotate JWT secret.
``revoke_refresh_tokens`` deletes the user's refresh tokens in the SAME
transaction: a separate delete could fail after the password commit and
leave a pre-change token still able to mint access tokens.
"""
from .hashing import hash_password
salt, pwd_hash = hash_password(new_password)
@ -563,6 +608,8 @@ def update_password(username: str, new_password: str) -> bool:
""",
(salt, pwd_hash, jwt_secret, username),
)
if revoke_refresh_tokens and cursor.rowcount > 0:
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.commit()
if cursor.rowcount > 0:
clear_bootstrap_password()

282
studio/backend/auth/terminal_prompt.py Normal file
View file

@ -0,0 +1,282 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Interactive terminal prompt that forces a bootstrap password change before
Studio is exposed on a public Cloudflare URL (``--secure`` / ``--cloudflare``).
Masked input echoes one ``*`` per keystroke (unlike ``getpass``). Works on
Windows (``msvcrt``) and Linux/macOS (``termios``). All output goes to stderr so
redirected stdout never swallows the prompt.
Mirrored for the CLI at ``unsloth_cli/commands/_password_prompt.py`` (the CLI
cannot import the Studio backend package); keep the two in sync.
"""
from __future__ import annotations
import os
import sys
from typing import Callable, TextIO
_CTRL_C = "\x03"
_CTRL_D = "\x04"
_CTRL_Z = "\x1a"
_BACKSPACES = ("\x7f", "\x08")
_SUBMITS = ("\r", "\n")
# Env var that supplies the initial admin password non-interactively (mirror in
# unsloth_cli/commands/_password_prompt.py). Keep the name in sync.
SUPPLIED_PASSWORD_ENV = "UNSLOTH_STUDIO_PASSWORD"
def _getch_windows() -> str: # pragma: no cover - exercised via fake on Linux CI
import msvcrt
ch = msvcrt.getwch()
# Function/arrow keys arrive as a two-wchar \x00/\xe0 sequence; consume the
# second half and report a no-op control char.
if ch in ("\x00", "\xe0"):
msvcrt.getwch()
return "\x00"
return ch
class _RestoreTtyOnSignals:
"""Restore terminal attrs if SIGTERM/SIGHUP kills the prompt mid-read.
A finally block can't run when a signal terminates the process, leaving the
shared terminal in cbreak/no-echo. Best-effort: no-op off the main thread or
where the signals are absent.
"""
def __init__(self, fd: int, old_attrs) -> None:
self._fd = fd
self._old_attrs = old_attrs
self._previous: list = []
def __enter__(self) -> "_RestoreTtyOnSignals":
import signal
import termios
def _restore_and_reraise(signum, frame):
termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs)
signal.signal(signum, signal.SIG_DFL)
signal.raise_signal(signum)
for name in ("SIGTERM", "SIGHUP"):
sig = getattr(signal, name, None)
if sig is None:
continue
try:
self._previous.append((sig, signal.signal(sig, _restore_and_reraise)))
except (ValueError, OSError): # non-main thread / unsupported
pass
return self
def __exit__(self, *exc) -> None:
import signal
for sig, previous in self._previous:
try:
signal.signal(sig, previous)
except (ValueError, OSError):
pass
class _prompt_raw_mode:
"""Hold cbreak + cleared ISIG (no echo) on stdin for the WHOLE prompt line,
restoring when the line finishes (and on SIGTERM/SIGHUP).
Echo must never re-enable mid-line: cbreak echoes on receipt, so a keystroke
arriving while echo is on would appear in cleartext. One cbreak block for the
whole line closes that window. No-op when stdin is not a real terminal, so
the _getch seam can be faked in tests.
"""
def __enter__(self) -> "_prompt_raw_mode":
self._fd = None
self._old_attrs = None
self._signals = None
try:
import termios
import tty
except ImportError: # non-POSIX (Windows uses msvcrt, no mode to hold)
return self
try:
fd = sys.stdin.fileno()
old_attrs = termios.tcgetattr(fd)
except (AttributeError, ValueError, OSError, termios.error):
return self # redirected / captured stdin (tests): nothing to hold
self._fd = fd
self._old_attrs = old_attrs
self._signals = _RestoreTtyOnSignals(fd, old_attrs)
self._signals.__enter__()
# cbreak (not raw) keeps output post-processing while disabling echo/line
# buffering. It leaves ISIG on, so clear it and surface Ctrl-C as \x03 to
# the caller loop, which restores the tty itself.
tty.setcbreak(fd, termios.TCSADRAIN)
new_attrs = termios.tcgetattr(fd)
new_attrs[3] &= ~termios.ISIG
termios.tcsetattr(fd, termios.TCSADRAIN, new_attrs)
return self
def __exit__(self, *exc) -> None:
if self._old_attrs is None:
return
import termios
try:
termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs)
finally:
if self._signals is not None:
self._signals.__exit__(*exc)
def _getch_posix() -> str: # pragma: no cover - needs a real tty
# Terminal already in cbreak+no-echo for the whole line (_prompt_raw_mode),
# so just read. Byte-at-a-time incremental decode so a multi-byte UTF-8 char
# straddling a read boundary isn't dropped.
import codecs
fd = sys.stdin.fileno()
decoder = codecs.getincrementaldecoder(sys.stdin.encoding or "utf-8")("replace")
while True:
b = os.read(fd, 1)
if not b:
return "" # stream EOF; caller raises EOFError
ch = decoder.decode(b)
if ch:
return ch
_getch: Callable[[], str] = _getch_windows if os.name == "nt" else _getch_posix
def _read_password(prompt: str, *, out: "TextIO | None" = None) -> str:
"""Read one masked line: echo ``*`` per char, support backspace editing.
Raises KeyboardInterrupt on Ctrl-C and EOFError on Ctrl-D/Ctrl-Z with an
empty buffer; the terminal is restored on every exit path.
"""
if out is None:
out = sys.stderr
out.write(prompt)
out.flush()
chars: list[str] = []
with _prompt_raw_mode():
while True:
key = _getch()
if key == "": # stream ended mid-line: abort, don't submit a partial
out.write("\n")
out.flush()
raise EOFError
for ch in key: # a paste can deliver several chars per read
if ch in _SUBMITS:
out.write("\n")
out.flush()
return "".join(chars)
if ch == _CTRL_C:
out.write("\n")
out.flush()
raise KeyboardInterrupt
if ch in (_CTRL_D, _CTRL_Z):
if not chars:
out.write("\n")
out.flush()
raise EOFError
continue # ignore mid-input
if ch in _BACKSPACES:
if chars:
chars.pop()
out.write("\b \b")
out.flush()
continue
if ch < " ": # other control characters (tab, escape, ...)
continue
chars.append(ch)
out.write("*")
out.flush()
def should_prompt_password_change(
*, tunnel_will_start: bool, requires_change: bool, stdin_isatty: bool, stderr_isatty: bool
) -> bool:
"""Whether to block startup on an interactive terminal password change.
True only when the tunnel is actually about to start, the admin still has
the seeded password, and both stdin and stderr are real terminals (headless
launches keep the bootstrap-timeout protection instead of hanging).
"""
return tunnel_will_start and requires_change and stdin_isatty and stderr_isatty
def prompt_for_password_change(
*,
min_length: int,
is_current_password: Callable[[str], bool],
apply_change: Callable[[str], None],
username: str = "unsloth",
out: "TextIO | None" = None,
) -> bool:
"""Force a new admin password before public exposure; True on success.
Loops until a valid, confirmed password is committed via ``apply_change``.
Ctrl-C / EOF returns False; the caller must then abort the launch.
"""
if out is None:
out = sys.stderr
out.write(
"\n"
"Unsloth Studio will be exposed on the public internet, so set a\n"
"password now. Ctrl+C to abort.\n\n"
)
out.flush()
try:
while True:
new_password = _read_password("New password: ", out = out)
if len(new_password) < min_length:
out.write(f"Password must be at least {min_length} characters; try again.\n")
out.flush()
continue
if is_current_password(new_password):
out.write(
"New password must differ from the current bootstrap password; try again.\n"
)
out.flush()
continue
confirmation = _read_password("Confirm new password: ", out = out)
if confirmation != new_password:
out.write("Passwords do not match; try again.\n")
out.flush()
continue
apply_change(new_password)
out.write(f"Password updated for '{username}'.\n")
out.flush()
return True
except (KeyboardInterrupt, EOFError):
out.write("Password change aborted; not exposing Studio.\n")
out.flush()
return False
def resolve_supplied_password(cli_value: "str | None", out: "TextIO | None" = None) -> "str | None":
"""Resolve a non-interactive initial admin password, or None if unset.
Precedence: an explicit ``--password`` (literal ``-`` reads a line from
stdin), then the ``UNSLOTH_STUDIO_PASSWORD`` env var; empty/omitted means off.
A literal argv value is visible in the process list, so a note points at the
env var or stdin instead. Mirror of the CLI helper -- keep the two in sync.
"""
if out is None:
out = sys.stderr
if cli_value == "-":
line = sys.stdin.readline()
if not line:
return None
return line.rstrip("\r\n") or None
if cli_value:
out.write(
"Note: --password is visible in the process list and shell history; "
f"prefer {SUPPLIED_PASSWORD_ENV} or --password - (stdin).\n"
)
out.flush()
return cli_value
return os.environ.get(SUPPLIED_PASSWORD_ENV) or None

4
studio/backend/colab.py
View file

@ -323,8 +323,8 @@ def start(port: int = 8888, *, cloudflare: bool = False):
logger.info(" Starting server...")
try:
# cloudflare=False: this helper owns the tunnel. run_server's default True
# would tunnel this 0.0.0.0 bind if Colab detection fails, breaking the opt-out.
# cloudflare=False: this helper owns the tunnel (Colab's own
# start(cloudflare=...) drives it), so pin it off explicitly.
app = run_server(
host = "0.0.0.0",
port = port,

7
studio/backend/core/data_recipe/service.py
View file

@ -9,6 +9,8 @@ import os
from pathlib import Path
from typing import Any
from utils.paths import recipe_datasets_root
from .jsonable import to_jsonable
from .local_callable_validators import (
register_oxc_local_callable_validators,
@ -277,6 +279,11 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None =
_apply_data_designer_image_context_patch()
from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports]
if artifact_path is None:
# DataDesigner defaults to cwd/artifacts; packaged Studio can run with
# cwd=/, so keep default callers on Studio's writable recipe artifact root.
artifact_path = str(recipe_datasets_root())
recipe = _strip_frontend_model_config_metadata(recipe)
model_providers = build_model_providers(recipe)
_validate_recipe_runtime_support(recipe, model_providers)

95
studio/backend/core/export/orchestrator.py
View file

@ -132,6 +132,11 @@ class ExportOrchestrator:
"""True while an export / load / cleanup command is running."""
return self._export_active
def is_worker_alive(self) -> bool:
"""True while the persistent export subprocess is running (op or idle)."""
proc = self._proc
return proc is not None and proc.is_alive()
def was_cancelled(self) -> bool:
"""True if the in-flight (or most recent) run was cancelled by the user."""
return self._cancel_requested
@ -204,6 +209,23 @@ class ExportOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new export subprocess."""
# Last-resort recheck for spawns outside an active op. Inside an op, _export_active is set and
# load_checkpoint already rechecked, so a reservation here is an install about to observe
# is_export_active() and abort; raising would kill this export for an install that never proceeds.
from utils.transformers_version import sidecar_swap_in_progress
from utils.transformers_version import sidecar_swap_kind
_swap_kind = sidecar_swap_kind()
# Inside an active op an INSTALL reservation is about to abort on the
# is_export_active check, but a lazy REPAIR has no such check and can be
# rebuilding the sidecar right now, so it must always refuse the spawn.
if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active):
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
"retry when it completes."
)
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
@ -231,11 +253,17 @@ class ExportOrchestrator:
adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep)
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the export subprocess."""
def _shutdown_subprocess(self, timeout: float = 10.0) -> bool:
"""Gracefully shut down the export subprocess.
Returns True only once the worker is confirmed dead. If it survives
terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives
SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the
pre-swap liveness guard can still observe the survivor instead of a cleared
handle and refuse the destructive sidecar swap."""
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
return True
self._drain_queue()
@ -265,10 +293,20 @@ class ExportOrchestrator:
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
# Survived SIGKILL (uninterruptible syscall): keep the handle so callers
# and the pre-swap guard see a live worker rather than a nulled one.
logger.error(
"Export subprocess still alive after terminate/kill; "
"preserving its handle for the pre-swap liveness check"
)
return False
self._proc = None
self._cmd_queue = None
self._resp_queue = None
logger.info("Export subprocess shut down")
return True
def _cleanup(self):
"""atexit handler."""
@ -339,9 +377,10 @@ class ExportOrchestrator:
if rtype == "status":
message = resp.get("message", "")
logger.info("Export subprocess status: %s", message)
# Surface status in the live log panel for high-level progress.
# One structured export_progress line per phase (consolidated in the
# server log, like training/download progress); also shown live.
if message:
logger.info("export_progress", phase = message)
self._append_log(
{
"stream": "status",
@ -409,14 +448,44 @@ class ExportOrchestrator:
self._export_active = True
op_success, op_message = False, ""
try:
# Handshake with the sidecar install route: _export_active is set above, so either this
# recheck refuses BEFORE tearing down the old worker (keeping the loaded checkpoint), or
# the install sees is_export_active() and 409s. The spawn-time recheck stays as a last resort.
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
from utils.transformers_version import SidecarSwapInProgress
op_message = (
"A transformers installation is replacing the latest "
"sidecar; retry when it completes."
)
raise SidecarSwapInProgress(op_message)
# Always kill any existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
if self._shutdown_subprocess() is False:
# Survivor still holds GPU memory (a wedged CUDA syscall outliving
# SIGKILL); its handle is kept so is_worker_alive() and the pre-swap
# guard still see it. Do not spawn a second worker over it -- fail so
# the load can retry once it exits.
op_message = (
"The current export worker did not exit and still holds GPU "
"memory; not starting a new checkpoint load over it. Retry shortly."
)
return False, op_message
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
self._spawn_subprocess(sub_config)
except Exception:
# The old worker is already gone; a stale current_checkpoint
# would make the Export page claim a loaded checkpoint that
# the next op then fails on with "no subprocess running".
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
raise
try:
resp = self._wait_response("loaded")
@ -560,6 +629,18 @@ class ExportOrchestrator:
self._export_active = True
op_success, op_message, op_output_path = False, "", None
try:
# Handshake with the sidecar install route (see load_checkpoint): _export_active is set
# above, so this recheck refuses before the command is sent, or the install sees the active
# op and 409s. Without it, an install would block in cleanup_memory behind a long export op.
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
from utils.transformers_version import SidecarSwapInProgress
op_message = (
"A transformers installation is replacing the latest "
"sidecar; retry when it completes."
)
raise SidecarSwapInProgress(op_message)
cmd = {"type": "export", "export_type": export_type, **params}
try:
self._send_cmd(cmd)

24
studio/backend/core/export/worker.py
View file

@ -236,6 +236,17 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path = cmd["checkpoint_path"]
max_seq_length = cmd.get("max_seq_length", 2048)
load_in_4bit = cmd.get("load_in_4bit", True)
# Latest-sidecar checkpoints load 16-bit here too: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (same flip as the chat worker).
if load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(checkpoint_path, cmd.get("hf_token")):
load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"export load (4-bit is disabled for brand-new architectures)",
checkpoint_path,
)
trust_remote_code = cmd.get("trust_remote_code", False)
# Auto-enable trust_remote_code for NemotronH/Nano models.
@ -387,6 +398,19 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
# orchestrator spawns a fresh subprocess per checkpoint load, resetting it.
_log_forward_gate.set()
# Phase milestone so the heavy export step shows in the server log; the
# merge/save/convert itself only forwards stdout to the live panel.
_phase = {
"merged": f"Exporting merged model ({cmd.get('format_type', '16-bit (FP16)')})...",
"gguf": f"Exporting GGUF ({cmd.get('quantization_method', 'Q4_K_M')})...",
"lora": "Exporting LoRA adapter...",
"base": "Exporting base model...",
}.get(export_type, f"Exporting ({export_type})...")
_send_response(
resp_queue,
{"type": "status", "message": _phase, "ts": time.time()},
)
output_path: Any = None
try:
if export_type == "merged":

387
studio/backend/core/inference/_html_to_md.py
View file

@ -7,6 +7,11 @@ Minimal HTML-to-Markdown converter using only the standard library.
Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line
``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic,
lists, tables, blockquotes, code blocks, and entity decoding.
``main_content=True`` also applies a readability-style heuristic: scope
conversion to the page's ``<article>`` (else ``<main>``) subtree when it
carries substantial text, and strip known boilerplate fragments (skip-links,
error placeholders, session banners, cookie prompts) from the result.
"""
from __future__ import annotations
@ -27,8 +32,138 @@ _SKIP_TAGS = frozenset(
"math",
"nav",
"footer",
# Never-rendered / form-chrome elements, not page content.
"template",
"dialog",
"button",
"select",
"datalist",
}
)
# <aside> is NOT skipped: docs use it for admonition callouts (real content);
# page-furniture asides are excluded by the main-content scoping pass instead.
# Void elements never produce an end tag, so they must not join the
# open-element stack used to bound hidden subtrees.
_VOID_TAGS = frozenset(
{
"area",
"base",
"br",
"col",
"embed",
"hr",
"img",
"input",
"link",
"meta",
"param",
"source",
"track",
"wbr",
}
)
def _style_hides_element(style: str) -> bool:
"""True when an inline ``style`` sets ``display:none`` / ``visibility:hidden``.
Parsed per property so an unrelated value that merely contains ``none`` is
not misread as hidden."""
lowered = style.lower()
if "none" not in lowered and "hidden" not in lowered:
return False
for declaration in style.split(";"):
prop, sep, value = declaration.partition(":")
if not sep:
continue
prop = prop.strip().lower()
# Drop any !important flag and keep the first token of the value.
value = value.split("!", 1)[0].strip().lower()
if prop == "display" and value == "none":
return True
if prop == "visibility" and value == "hidden":
return True
return False
def _is_hidden_element(attr_dict: dict) -> bool:
"""True when the element is not rendered: ``hidden`` attribute,
``aria-hidden="true"``, or an inline ``style`` hiding it. Such JS-only
placeholders ship in the HTML but must not reach the output. ``hidden`` is
enumerated: any present value (even ``hidden="false"``) means not rendered."""
if "hidden" in attr_dict:
return True
if (attr_dict.get("aria-hidden") or "").strip().lower() == "true":
return True
return _style_hides_element(attr_dict.get("style") or "")
# HTML5 optional end tags: a listed start tag implicitly closes an open element
# of the key type (as browsers do), else an unclosed ``<p hidden>``/``<li hidden>``
# swallows every following sibling. Keys: closable elements; values: closers.
_P_CLOSING_TAGS = frozenset(
{
"address",
"article",
"aside",
"blockquote",
"details",
"div",
"dl",
"fieldset",
"figcaption",
"figure",
"footer",
"form",
"h1",
"h2",
"h3",
"h4",
"h5",
"h6",
"header",
"hgroup",
"hr",
"main",
"menu",
"nav",
"ol",
"p",
"pre",
"section",
"table",
"ul",
}
)
_IMPLICIT_CLOSERS: dict = {
"p": _P_CLOSING_TAGS,
"li": frozenset({"li"}),
"dt": frozenset({"dt", "dd"}),
"dd": frozenset({"dt", "dd"}),
"tr": frozenset({"tr"}),
"td": frozenset({"td", "th", "tr"}),
"th": frozenset({"td", "th", "tr"}),
"option": frozenset({"option", "optgroup"}),
"optgroup": frozenset({"optgroup"}),
}
# Item tag -> container tags that re-scope it: a nested container makes an inner
# item a descendant, not an optional-close sibling, so recovery must stop there
# rather than close (and un-hide) the outer item and leak its nested content.
_CLOSE_BARRIERS: dict = {
"li": frozenset({"ul", "ol", "menu"}),
"dt": frozenset({"dl"}),
"dd": frozenset({"dl"}),
"tr": frozenset({"table"}),
"td": frozenset({"table"}),
"th": frozenset({"table"}),
"option": frozenset({"select", "datalist"}),
"optgroup": frozenset({"select", "datalist"}),
}
_BLOCK_TAGS = frozenset(
{
"p",
@ -51,13 +186,33 @@ _INLINE_EMPHASIS = {"strong": "**", "b": "**", "em": "*", "i": "*"}
class _MarkdownRenderer(HTMLParser):
"""HTMLParser subclass that emits Markdown tokens into a list."""
"""HTMLParser subclass that emits Markdown tokens into a list.
def __init__(self):
``scope_tags`` restricts emission to the subtree(s) of the given tags
(e.g. ``{"article"}``): outside them every handler is a no-op, which is
how the readability-style main-content pass drops page furniture.
"""
def __init__(self, scope_tags: frozenset[str] | None = None):
super().__init__(convert_charrefs = False)
self._out: list[str] = []
self._skip_depth: int = 0
# Main-content scoping: emit only while inside a scope tag.
self._scope_tags = scope_tags
self._scope_depth: int = 0
# Output boundaries per top-level scope element, so a caller can size each
# candidate alone and a swarm of tiny sibling cards can't clear the threshold.
self.scope_segments: list[str] = []
self._scope_seg_start: int | None = None
# Hidden-subtree tracking: stack of open non-void tags plus the indices
# where a hidden element started. End tags pop to the matching tag, so
# an omitted </p>/<li> close cannot leave the renderer stuck hidden.
self._open_tags: list[str] = []
self._hidden_marks: list[int] = []
# Link state
self._link_href: str | None = None
self._link_text_parts: list[str] = []
@ -150,16 +305,95 @@ class _MarkdownRenderer(HTMLParser):
# ------------------------------------------------------------------
# Tag handlers
# ------------------------------------------------------------------
# Structural bookkeeping shared by every start tag (skip/hidden/scope).
def _close_implicit(self, tag: str) -> None:
"""HTML5 optional-end-tag recovery for a start tag about to open.
Pops each implicitly-closed ancestor (and its hidden marks), scanning the
whole stack so an open ``<p>``/``<li>`` still closes under an unclosed inline
``<span>``. Stops at a ``_CLOSE_BARRIERS`` container so recovery never crosses
a nested list/table/dl and leaks the outer item's hidden content. Runs even
for skipped ``<nav>``/``<footer>``, which also close ``<p>``."""
barriers = _CLOSE_BARRIERS.get(tag, ())
while True:
close_at = None
for i in range(len(self._open_tags) - 1, -1, -1):
name = self._open_tags[i]
if tag in _IMPLICIT_CLOSERS.get(name, ()):
close_at = i
break
# A barrier container re-scopes the item; stop before it.
if name in barriers:
break
if close_at is None:
break
del self._open_tags[close_at:]
while self._hidden_marks and self._hidden_marks[-1] >= close_at:
self._hidden_marks.pop()
def _enter_tag(self, tag: str, attr_dict: dict) -> bool:
"""Track open/hidden/scope state; return True when the tag's content
should be rendered (False = suppressed). Caller runs ``_close_implicit``
first so recovery also fires for skipped tags."""
if tag not in _VOID_TAGS:
self._open_tags.append(tag)
if _is_hidden_element(attr_dict):
self._hidden_marks.append(len(self._open_tags) - 1)
elif _is_hidden_element(attr_dict):
# Void elements never join the stack, so suppress a hidden one inline.
return False
if self._scope_tags is not None and tag in self._scope_tags:
if self._scope_depth == 0:
self._scope_seg_start = len(self._out)
self._scope_depth += 1
if self._hidden_marks:
return False
if self._scope_tags is not None and self._scope_depth == 0:
return False
return True
def _exit_tag(self, tag: str) -> bool:
"""Pop to the matching open tag; return True when the end tag should
be rendered (False = it closed inside a hidden / out-of-scope region)."""
suppressed = bool(self._hidden_marks) or (
self._scope_tags is not None and self._scope_depth == 0
)
if tag not in _VOID_TAGS:
# Pop to the innermost matching open tag (recovers omitted closes).
for i in range(len(self._open_tags) - 1, -1, -1):
if self._open_tags[i] == tag:
del self._open_tags[i:]
while self._hidden_marks and self._hidden_marks[-1] >= i:
self._hidden_marks.pop()
break
if self._scope_tags is not None and tag in self._scope_tags and self._scope_depth > 0:
self._scope_depth -= 1
if self._scope_depth == 0 and self._scope_seg_start is not None:
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
self._scope_seg_start = None
return not suppressed
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
tag = tag.lower()
if self._skip_depth:
# Inside a skipped subtree: only track nested skip depth.
if tag in _SKIP_TAGS:
self._skip_depth += 1
return
# Recover optional end tags before the skip decision: a skipped
# <nav>/<footer> still implicitly closes an open <p>, releasing its
# hidden mark so following siblings render.
self._close_implicit(tag)
if tag in _SKIP_TAGS:
self._skip_depth += 1
return
if self._skip_depth:
return
attr_dict = dict(attrs)
if not self._enter_tag(tag, attr_dict):
return
if tag in _HEADING_TAGS:
level = int(tag[1])
@ -250,6 +484,9 @@ class _MarkdownRenderer(HTMLParser):
if self._skip_depth:
return
if not self._exit_tag(tag):
return
if tag in _HEADING_TAGS:
self._emit("\n\n")
@ -308,8 +545,13 @@ class _MarkdownRenderer(HTMLParser):
# ------------------------------------------------------------------
# Text / entity handlers
# ------------------------------------------------------------------
def _text_suppressed(self) -> bool:
if self._skip_depth or self._hidden_marks:
return True
return self._scope_tags is not None and self._scope_depth == 0
def handle_data(self, data: str) -> None:
if self._skip_depth:
if self._text_suppressed():
return
if self._in_pre:
self._pre_parts.append(data)
@ -326,12 +568,12 @@ class _MarkdownRenderer(HTMLParser):
self._emit(text)
def handle_entityref(self, name: str) -> None:
if self._skip_depth:
if self._text_suppressed():
return
self._emit(html.unescape(f"&{name};"))
def handle_charref(self, name: str) -> None:
if self._skip_depth:
if self._text_suppressed():
return
self._emit(html.unescape(f"&#{name};"))
@ -366,6 +608,14 @@ class _MarkdownRenderer(HTMLParser):
else:
self._out.append("\n\n" + prefixed + "\n\n")
# A scope left open by truncated HTML never reached _exit_tag, so its output
# never joined scope_segments and would score 0. Flush the still-open segment
# here (after the side-buffers) so a truncated main-content page is scored.
if self._scope_seg_start is not None:
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
self._scope_seg_start = None
self._scope_depth = 0
# Post-processing
def _cleanup(text: str) -> str:
@ -399,17 +649,124 @@ def _cleanup(text: str) -> str:
return "\n".join(out).strip()
# Public API
def html_to_markdown(source_html: str) -> str:
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
# Known boilerplate fragments stripped from main-content conversions, matched
# only against short lines. Sources: GitHub page furniture / client-side error
# placeholders, skip-links, cookie banners.
_BOILERPLATE_FRAGMENTS = (
"skip to content",
"skip to main content",
"there was an error while loading",
"please reload this page",
"you can't perform that action at this time",
"you signed in with another tab or window",
"you signed out in another tab or window",
"you switched accounts on another tab or window",
"reload to refresh your session",
"you must be signed in to change notification settings",
"uh oh!",
"{{ message }}",
"this website uses cookies",
"we use cookies",
"accept all cookies",
"manage cookie preferences",
)
# Only shorter lines are eligible for boilerplate dropping; real content
# sentences quoting a fragment run longer.
_BOILERPLATE_MAX_LINE_CHARS = 300
``<script>``, ``<style>``, and ``<head>`` are stripped entirely.
"""
# Normalize line endings before parsing.
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
renderer = _MarkdownRenderer()
# Normalized furniture phrases for whole-segment matching. See _line_is_boilerplate.
_BOILERPLATE_NORMALIZED = frozenset(
re.sub(r"\s+", " ", fragment).strip().casefold().rstrip(".!:")
for fragment in _BOILERPLATE_FRAGMENTS
)
def _line_is_boilerplate(line: str) -> bool:
"""True only when a whole line is composed of known furniture phrases.
Splits on sentence terminators and requires every segment to be furniture, so a
line stacking several phrases is dropped while prose that merely quotes one is
kept (its other words leave a non-furniture segment)."""
normalized = re.sub(r"\s+", " ", line).strip().casefold()
if not normalized:
return False
segments = [segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized)]
segments = [segment for segment in segments if segment]
return bool(segments) and all(segment in _BOILERPLATE_NORMALIZED for segment in segments)
def _strip_boilerplate_lines(text: str) -> str:
"""Drop short lines that consist entirely of known page-furniture phrases.
Fenced code blocks are preserved verbatim: boilerplate never renders
inside ``<pre>``, while READMEs legitimately quote error strings."""
out: list[str] = []
in_fence = False
for line in text.split("\n"):
if line.lstrip().startswith("```"):
in_fence = not in_fence
out.append(line)
continue
if not in_fence and len(line) <= _BOILERPLATE_MAX_LINE_CHARS and _line_is_boilerplate(line):
continue
out.append(line)
# Collapse blank runs the dropped lines may have left behind.
return re.sub(r"\n{3,}", "\n\n", "\n".join(out)).strip()
def _render(source_html: str, scope_tags: frozenset[str] | None) -> str:
renderer = _MarkdownRenderer(scope_tags = scope_tags)
renderer.feed(source_html)
renderer.close()
renderer.flush_pending()
raw = "".join(renderer._out)
return _cleanup(raw)
def _select_main_scope_render(source_html: str, tag: str) -> tuple[int, str]:
"""Length and boilerplate-stripped render of the largest single ``<tag>``
subtree. Sizing candidates one at a time stops many tiny sibling cards from
clearing the threshold together, and returning that one subtree keeps
unrelated siblings (related cards, comment threads) out of the output."""
renderer = _MarkdownRenderer(scope_tags = frozenset({tag}))
renderer.feed(source_html)
renderer.close()
renderer.flush_pending()
best_len = 0
best_render = ""
for seg in renderer.scope_segments:
rendered = _strip_boilerplate_lines(_cleanup(seg))
if len(rendered) > best_len:
best_len = len(rendered)
best_render = rendered
return best_len, best_render
# A scoped conversion below this size is judged not to be the page's main
# content (e.g. an empty <article> stub) and the next candidate is tried.
_MIN_MAIN_CONTENT_CHARS = 200
# Public API
def html_to_markdown(source_html: str, *, main_content: bool = False) -> str:
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
``<script>``, ``<style>``, and ``<head>`` are stripped entirely, as are
subtrees hidden from rendering (``hidden`` / ``aria-hidden="true"``).
``main_content=True`` applies a readability-style heuristic for page
fetches: prefer the ``<article>`` subtree (GitHub renders READMEs there),
then ``<main>``, falling back to the whole document, and strip known
boilerplate fragments from the result.
"""
# Normalize line endings before parsing.
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
if main_content:
for scope_tag in ("article", "main"):
# Render only the chosen subtree so sibling <article>/<main>
# elements do not leak in once the largest passes the size gate.
length, rendered = _select_main_scope_render(source_html, scope_tag)
if length >= _MIN_MAIN_CONTENT_CHARS:
return rendered
return _strip_boilerplate_lines(_render(source_html, None))
return _render(source_html, None)

110
studio/backend/core/inference/_vulkan_probe.py Normal file
View file

@ -0,0 +1,110 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Standalone free-VRAM probe for the bundled ggml Vulkan backend.
Run in a short-lived subprocess (``python _vulkan_probe.py <bindir>``) so the
Vulkan instance never lives in the long-running backend process. Loads the
bundled ggml Vulkan backend from ``<bindir>`` and prints one
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` line per device to stdout.
Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
Uses only the standard library so it stays runnable as a bare script.
"""
import ctypes
import os
import sys
# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ...
_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
def _igpu_flags(base, lib, count: int) -> list[bool]:
"""Per-device integrated-GPU flags via ggml's backend registry.
The Vulkan reg enumerates devices in the same order as
``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
i``), so reg index == device ordinal. Returns all-False on any failure so
the reader never over-caps a discrete card.
"""
flags = [False] * count
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t
base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p]
base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p
base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
base.ggml_backend_dev_type.restype = ctypes.c_int
base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p]
reg = lib.ggml_backend_vk_reg()
if not reg:
return flags
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
except Exception:
# Best-effort: any failure degrades to "discrete" so the memory
# readings still get through instead of crashing the probe.
pass
return flags
def main() -> int:
if len(sys.argv) < 2:
return 0
bindir = sys.argv[1]
# Hold add_dll_directory's handle for the rest of main() (the documented
# idiom) so bindir stays on the search path while the sibling ggml DLLs
# resolve below.
_dll_dir = None
if sys.platform == "win32":
base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll"
try:
_dll_dir = os.add_dll_directory(bindir)
except Exception:
pass
else:
base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
# RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr
# falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode).
_rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0)
try:
base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global)
lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global)
except OSError as e:
print(f"ggml-vulkan load failed: {e}", file = sys.stderr)
return 1
lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int
lib.ggml_backend_vk_get_device_count.argtypes = []
lib.ggml_backend_vk_get_device_memory.restype = None
lib.ggml_backend_vk_get_device_memory.argtypes = [
ctypes.c_int,
ctypes.POINTER(ctypes.c_size_t),
ctypes.POINTER(ctypes.c_size_t),
]
count = lib.ggml_backend_vk_get_device_count()
igpu = _igpu_flags(base, lib, count)
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
sys.stdout.write("\n".join(rows))
return 0
if __name__ == "__main__":
raise SystemExit(main())

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

@ -12,6 +12,43 @@ import json
import logging
from typing import Optional
_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"
def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str:
"""Return the trailing open ``<think>`` prefill of a rendered prompt.
Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation
prompt with ``<think>\\n`` so the model starts reasoning immediately.
Because that opening tag is part of the *prompt*, skip_prompt streaming
never emits it, and the frontend's ``<think>``/``</think>`` parser shows
the reasoning as plain text instead of a thinking block. (The GGUF path
is unaffected: llama-server's reasoning parser returns
``reasoning_content``, which gets re-wrapped in think tags.)
Returns the exact prompt tail to re-emit at the start of the generated
stream (e.g. ``"<think>\\n"``), or ``""`` when the prompt does not end
with an open think block, including the ``enable_thinking=False`` case
where templates prefill an already-closed ``<think>\\n\\n</think>``.
``special_tokens`` is the tokenizer's special-token list. If ``</think>``
is one, the streamer's skip_special_tokens strips the model's closing tag,
so re-emitting the open would leave an unclosed block that swallows the
answer. In that case return ``""`` and fall back to plain text.
"""
if not prompt:
return ""
open_idx = prompt.rfind(_THINK_OPEN)
if open_idx == -1:
return ""
tail = prompt[open_idx:]
if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN:
return ""
if special_tokens and _THINK_CLOSE in set(special_tokens):
return ""
return tail
logger = logging.getLogger(__name__)

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

@ -15,6 +15,7 @@ from pathlib import Path
from typing import Optional, Union, Generator, Tuple
from utils.models import ModelConfig, get_base_model_from_lora
from utils.paths import is_model_cached
from utils.transformers_dtype import dtype_kwargs
from utils.utils import format_error_message
from utils.hardware import (
get_device,
@ -440,7 +441,7 @@ class InferenceBackend:
feature_extractor = tokenizer.feature_extractor,
processor = tokenizer,
return_language = True,
torch_dtype = torch.float16,
**dtype_kwargs(torch.float16),
)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
@ -832,6 +833,7 @@ class InferenceBackend:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
presence_penalty: float = 0.0,
):
@ -885,6 +887,7 @@ class InferenceBackend:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
)
@ -1178,13 +1181,22 @@ class InferenceBackend:
add_special_tokens = False,
return_tensors = "pt",
).to(model.device)
prompt_text = input_text
else:
# Text-only path for a vision model
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device)
prompt_text = formatted_prompt
# Stream with TextIteratorStreamer + background thread
try:
from core.inference.chat_template_helpers import detect_think_prefill
# Re-emit an open <think> prefill swallowed by skip_prompt (see
# generate_stream).
think_prefix = detect_think_prefill(
prompt_text, getattr(raw_tokenizer, "all_special_tokens", None)
)
from transformers import TextIteratorStreamer
import threading
@ -1233,7 +1245,11 @@ class InferenceBackend:
thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
output = think_prefix
# Emit the prefilled <think> before the first token so the block
# renders during prompt prefill (which can take seconds).
if think_prefix:
yield think_prefix
from queue import Empty
generation_complete = False
@ -1467,6 +1483,16 @@ class InferenceBackend:
from transformers import TextIteratorStreamer
import threading
from core.inference.chat_template_helpers import detect_think_prefill
# skip_prompt swallows an open <think> prefilled by the template;
# re-emit it so the frontend can render the thinking block.
# gpt-oss emits its own tags via HarmonyTextStreamer.
think_prefix = (
""
if self._is_gpt_oss_model()
else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None))
)
# gpt-oss models: HarmonyTextStreamer parses the multi-channel
# harmony protocol into <think> tags
@ -1550,7 +1576,11 @@ class InferenceBackend:
thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
output = think_prefix
# Emit the prefilled <think> before the first token so the block
# renders during prompt prefill (which can take seconds).
if think_prefix:
yield think_prefix
from queue import Empty
generation_complete = False

368
studio/backend/core/inference/llama_admission.py Normal file
View file

@ -0,0 +1,368 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Admission control for local llama-server generation requests.
The helpers in this module deliberately know nothing about FastAPI, SSE, or the
OpenAI-compatible route shape. They only coordinate how many upstream generation
requests may be active for one llama-server backend and provide a cancellable
FIFO queue for excess requests.
"""
from __future__ import annotations
import asyncio
import os
import threading
from collections import deque
from dataclasses import dataclass
from typing import Deque, Optional
ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL"
ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT"
ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL"
ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE"
DEFAULT_ADMISSION_ENABLED = True
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0
DEFAULT_ADMISSION_MAX_QUEUE = 64
@dataclass(frozen = True)
class LlamaAdmissionConfig:
enabled: bool = DEFAULT_ADMISSION_ENABLED
queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE
@dataclass(frozen = True)
class LlamaAdmissionSnapshot:
key: str
capacity: int
active: int
queued: int
class LlamaAdmissionError(Exception):
def __init__(
self,
message: str,
*,
snapshot: Optional[LlamaAdmissionSnapshot] = None,
):
super().__init__(message)
self.snapshot = snapshot
class LlamaAdmissionQueueFull(LlamaAdmissionError):
pass
class LlamaAdmissionTimeout(LlamaAdmissionError):
pass
class LlamaAdmissionCancelled(LlamaAdmissionError):
pass
def _bool_env(name: str, default: bool) -> bool:
value = os.environ.get(name)
if value is None or not value.strip():
return default
value = value.strip().lower()
if value in {"1", "true", "yes", "on"}:
return True
if value in {"0", "false", "no", "off"}:
return False
return default
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = float(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else None
def _positive_float_env(name: str, default: float) -> float:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = float(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else default
def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = int(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else None
def llama_admission_config_from_env() -> LlamaAdmissionConfig:
return LlamaAdmissionConfig(
enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED),
queue_timeout_s = _optional_positive_float_env(
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S,
),
keepalive_interval_s = _positive_float_env(
ADMISSION_KEEPALIVE_INTERVAL_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
),
max_queue = _optional_positive_int_env(
ADMISSION_MAX_QUEUE_ENV,
DEFAULT_ADMISSION_MAX_QUEUE,
),
)
@dataclass
class _Waiter:
loop: asyncio.AbstractEventLoop
future: asyncio.Future
cancelled: bool = False
granted_lease: Optional["LlamaAdmissionLease"] = None
class LlamaAdmissionLease:
def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
self._queue = queue
self._released = False
self._release_lock = threading.Lock()
def release(self) -> None:
queue = None
with self._release_lock:
if self._released:
return
self._released = True
queue = self._queue
if queue is not None:
queue.release()
async def __aenter__(self) -> "LlamaAdmissionLease":
return self
async def __aexit__(self, *_args) -> None:
self.release()
class LlamaAdmissionReservation:
def __init__(
self,
*,
queue: Optional["LlamaAdmissionQueue"],
lease: Optional[LlamaAdmissionLease] = None,
waiter: Optional[_Waiter] = None,
snapshot: Optional[LlamaAdmissionSnapshot] = None,
):
self._queue = queue
self._lease = lease
self._waiter = waiter
self.snapshot = snapshot
@property
def is_cancelled(self) -> bool:
return self._lease is None and self._waiter is None
def lease_nowait(self) -> Optional[LlamaAdmissionLease]:
if self._lease is not None:
return self._lease
if self._waiter is None or not self._waiter.future.done():
return None
if self._waiter.future.cancelled():
self._waiter.cancelled = True
self._waiter = None
return None
self._lease = self._waiter.future.result()
self._waiter = None
return self._lease
async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]:
lease = self.lease_nowait()
if lease is not None:
return lease
if self._waiter is None:
return None
waiter = self._waiter
try:
await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s)
except asyncio.CancelledError:
if waiter.future.cancelled():
waiter.cancelled = True
if self._waiter is waiter:
self._waiter = None
return None
raise
return self.lease_nowait()
def cancel(self) -> None:
lease = self.lease_nowait()
if lease is not None:
lease.release()
self._lease = None
return
if self._queue is not None and self._waiter is not None:
self._queue.cancel(self._waiter)
self._waiter = None
def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]:
if self._queue is None:
return self.snapshot
return self._queue.snapshot()
class LlamaAdmissionQueue:
def __init__(self, key: str):
self.key = key
self._lock = threading.Lock()
self._active = 0
self._capacity = 1
self._waiters: Deque[_Waiter] = deque()
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
capacity = max(1, int(capacity or 1))
if not config.enabled:
return LlamaAdmissionReservation(
queue = None,
lease = LlamaAdmissionLease(None),
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0),
)
loop = asyncio.get_running_loop()
with self._lock:
self._capacity = capacity
self._prune_waiters_locked()
self._grant_waiters_locked()
if self._active < self._capacity and not self._waiters:
self._active += 1
return LlamaAdmissionReservation(
queue = self,
lease = LlamaAdmissionLease(self),
snapshot = self._snapshot_locked(),
)
if config.max_queue is not None and len(self._waiters) >= config.max_queue:
raise LlamaAdmissionQueueFull(
"llama-server generation queue is full",
snapshot = self._snapshot_locked(),
)
waiter = _Waiter(
loop = loop,
future = loop.create_future(),
)
self._waiters.append(waiter)
return LlamaAdmissionReservation(
queue = self,
waiter = waiter,
snapshot = self._snapshot_locked(),
)
def release(self) -> None:
with self._lock:
if self._active > 0:
self._active -= 1
self._grant_waiters_locked()
def cancel(self, waiter: _Waiter) -> None:
lease_to_release = None
with self._lock:
waiter.cancelled = True
try:
self._waiters.remove(waiter)
except ValueError:
pass
if waiter.granted_lease is not None:
lease_to_release = waiter.granted_lease
waiter.granted_lease = None
if not waiter.future.done():
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
if lease_to_release is not None:
lease_to_release.release()
def snapshot(self) -> LlamaAdmissionSnapshot:
with self._lock:
self._prune_waiters_locked()
return self._snapshot_locked()
def is_idle(self) -> bool:
with self._lock:
self._prune_waiters_locked()
return self._active == 0 and not self._waiters
def _grant_waiters_locked(self) -> None:
self._prune_waiters_locked()
while self._waiters and self._active < self._capacity:
waiter = self._waiters.popleft()
if waiter.cancelled or waiter.future.done():
continue
self._active += 1
lease = LlamaAdmissionLease(self)
waiter.granted_lease = lease
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None:
if waiter.cancelled or waiter.future.done():
waiter.granted_lease = None
if not waiter.future.done():
waiter.future.cancel()
lease.release()
return
try:
waiter.future.set_result(lease)
waiter.granted_lease = None
except asyncio.InvalidStateError:
waiter.granted_lease = None
lease.release()
def _prune_waiters_locked(self) -> None:
self._waiters = deque(
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
)
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:
return LlamaAdmissionSnapshot(
key = self.key,
capacity = self._capacity,
active = self._active,
queued = len(self._waiters),
)
_QUEUES_LOCK = threading.Lock()
_QUEUES: dict[str, LlamaAdmissionQueue] = {}
def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue:
with _QUEUES_LOCK:
queue = _QUEUES.get(key)
if queue is None:
queue = LlamaAdmissionQueue(key)
_QUEUES[key] = queue
# base_url carries a fresh ephemeral port on every model load, so
# each load registers a new key. Drop the now-idle queues from prior
# loads so the registry can't grow without bound on a long-running
# server. Queues with in-flight requests are kept until they drain.
for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]:
del _QUEUES[stale_key]
return queue
def reset_llama_admission_queues() -> None:
with _QUEUES_LOCK:
_QUEUES.clear()

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

File diff suppressed because it is too large Load diff

774
studio/backend/core/inference/mcp_client.py
View file

@ -4,11 +4,16 @@
from __future__ import annotations
import asyncio
import atexit
import concurrent.futures
import hashlib
import json
import os
import shlex
import sys
import threading
import time
import uuid
from typing import Any, Optional
from loggers import get_logger
@ -115,6 +120,18 @@ def join_stdio_command(parts: list[str]) -> str:
return shlex.join(parts)
def _stdio_log_id(url: str) -> str:
"""A non-secret label for logs. stdio commands can embed credentials in argv
(e.g. ``npx server --token sk-...``), so never log the raw command; use the
executable basename plus a short digest of the full command instead."""
try:
parts = parse_stdio_command(url)
exe = os.path.basename(parts[0]) if parts else "<empty>"
except Exception: # noqa: BLE001
exe = "<invalid>"
return f"{exe}#{hashlib.sha256(url.encode()).hexdigest()[:12]}"
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (bypassing the
sandbox), so allowed only when the host is the user's own machine. On startup
@ -192,7 +209,6 @@ async def clear_oauth_tokens_async(url: str) -> None:
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
await auth.token_storage_adapter.clear()
except Exception as exc: # noqa: BLE001
# Cleanup is best-effort; the row delete still wins.
logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
@ -237,6 +253,502 @@ def _client(
return Client(transport_cls(url = url, headers = headers or None, auth = auth))
# Persistent stdio sessions: a stdio MCP server owns live state (a browser, a
# DB handle), so keep one connected client per (command, env, chat session) on
# a dedicated event-loop thread instead of respawning per call.
_STDIO_SESSION_IDLE_TTL = 300.0
_STDIO_SESSION_REAP_INTERVAL = 30.0
_STDIO_CONNECT_TIMEOUT = 60.0 # allows first-run `npx -y ...` package download
_STDIO_CLOSE_TIMEOUT = 10.0
_STDIO_WEDGE_MARGIN = 15.0
# Cap concurrent persistent sessions: each owns a subprocess + loop thread, and
# the scope includes a caller-supplied thread_id, so an unbounded cache is a
# resource-exhaustion surface. Overridable via env for large deployments.
try:
_STDIO_MAX_SESSIONS = max(1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32")))
except ValueError:
_STDIO_MAX_SESSIONS = 32
def _is_tool_error(exc: BaseException) -> bool:
"""A tool-level failure (the tool ran and errored) leaves the transport alive,
so the session is kept; fastmcp raises ToolError for these. Anything else from
call_tool is transport-level. Version-safe (fastmcp 3.0.2 has no dead probe)."""
try:
from fastmcp.exceptions import ToolError
except Exception: # noqa: BLE001
return False
return isinstance(exc, ToolError)
def _transport_dead(session) -> bool:
"""Best-effort, version-adaptive liveness probe for a cached stdio client.
``Client.is_connected()`` only checks a session object exists, not that the
subprocess is alive, so it is never used here. Returns True only when the
transport is positively gone; unknown returns False (the call surfaces it)."""
client = getattr(session, "client", None)
if client is None:
return True
transport = getattr(client, "transport", None)
probe = getattr(transport, "_is_session_dead", None)
if callable(probe):
try:
if probe():
return True
except Exception: # noqa: BLE001
pass
connect_task = getattr(transport, "_connect_task", None)
if connect_task is not None:
try:
if connect_task.done():
return True
except Exception: # noqa: BLE001
pass
return False
class _SessionWedged(Exception):
pass
class _SessionClosed(Exception):
"""The session was closed (server update/delete/shutdown) mid-call."""
def _abort_future(future) -> None:
# Let the cancelled coroutine unwind before its loop is stopped.
future.cancel()
try:
future.result(1.0)
except BaseException: # noqa: BLE001
pass
class _StdioSession:
def __init__(self, url: str, headers: Optional[dict]):
self.url = url
self.headers = headers
self.client = None
self.closed = threading.Event()
self.defunct = False # discarded; close once in_flight drains (see _retire)
self._close_lock = threading.Lock()
self.call_lock = threading.Lock() # serializes tool calls on this session
self.last_used = time.monotonic()
self.in_flight = 0 # guarded by _stdio_sessions_lock
# On Windows a bare new_event_loop() can be a SelectorEventLoop (if any
# component set that policy), which cannot spawn subprocesses natively;
# force a ProactorEventLoop so the stdio transport always works.
if sys.platform == "win32":
self.loop = asyncio.ProactorEventLoop()
else:
self.loop = asyncio.new_event_loop()
self._thread = threading.Thread(
target = self._run_loop, name = "mcp-stdio-session", daemon = True
)
self._thread.start()
def _run_loop(self) -> None:
asyncio.set_event_loop(self.loop)
try:
self.loop.run_forever()
finally:
self.loop.close()
def connect(self, timeout: Optional[float], cancel_event) -> None:
async def _open():
client = _client(self.url, self.headers)
await client.__aenter__()
# Publish on the loop thread with no await in between: if an abort
# races a just-completed connect, close() still sees the client and
# __aexit__s it instead of orphaning the subprocess.
self.client = client
return client
future = asyncio.run_coroutine_threadsafe(_open(), self.loop)
# timeout=None means unlimited (no connect deadline); a finite caller
# timeout still bounds connect by min(timeout, _STDIO_CONNECT_TIMEOUT).
window = None if timeout is None else min(timeout, _STDIO_CONNECT_TIMEOUT)
deadline = None if window is None else time.monotonic() + window
while True:
if cancel_event is not None and cancel_event.is_set():
_abort_future(future)
raise _MCPCancelled
try:
future.result(0.05)
return
except (concurrent.futures.TimeoutError, asyncio.TimeoutError):
if future.done():
raise # the connect itself failed fast; don't wait out the window
if deadline is not None and time.monotonic() >= deadline:
_abort_future(future)
raise asyncio.TimeoutError
def is_connected(self) -> bool:
client = self.client
if client is None:
return False
probe = getattr(client, "is_connected", None)
try:
return bool(probe()) if callable(probe) else True
except Exception:
return False
def run(self, coro, timeout: Optional[float]):
self.last_used = time.monotonic()
future = asyncio.run_coroutine_threadsafe(coro, self.loop)
# The coroutine enforces the tool timeout; the margin only catches a
# wedged loop. No deadline at all when the caller set none -- but poll
# so a session closed under us (server update/delete) can't hang the
# request thread forever on a stopped loop.
deadline = None if timeout is None else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN
try:
while True:
try:
return future.result(0.25)
except concurrent.futures.CancelledError:
# Only close() cancels in-flight tasks (in _shutdown).
raise _SessionClosed
except (concurrent.futures.TimeoutError, asyncio.TimeoutError):
if future.done():
raise # the call's own timeout; the session stays usable
if self.closed.is_set():
future.cancel()
raise _SessionClosed
if deadline is not None and time.monotonic() >= deadline:
future.cancel()
raise _SessionWedged
finally:
self.last_used = time.monotonic()
def close(self) -> None:
# Idempotent: a discard racing close_stdio_sessions() may close twice.
# Setting `closed` first also unblocks run() waiters (they poll it).
with self._close_lock:
if self.closed.is_set():
return
self.closed.set()
loop = getattr(self, "loop", None)
loop_alive = loop is not None and not loop.is_closed()
if loop_alive:
async def _shutdown() -> None:
# Runs on the loop thread, so it serializes with an aborted
# connect() that finished anyway and just published its client.
client, self.client = self.client, None
if client is not None:
await client.__aexit__(None, None, None)
# Cancel in-flight calls so they unwind before loop.stop
# (their run() waiters have already been released via `closed`).
for task in asyncio.all_tasks():
if task is not asyncio.current_task():
task.cancel()
try:
asyncio.run_coroutine_threadsafe(_shutdown(), loop).result(_STDIO_CLOSE_TIMEOUT)
except Exception as exc: # noqa: BLE001
logger.warning(
"MCP stdio session close failed for %s: %s",
_stdio_log_id(getattr(self, "url", "")),
exc,
)
try:
loop.call_soon_threadsafe(loop.stop)
except RuntimeError:
pass
else:
self.client = None
thread = getattr(self, "_thread", None)
if thread is not None:
thread.join(timeout = 5.0)
_stdio_sessions: dict[tuple, _StdioSession] = {}
# Per-key locks so a slow connect/close never blocks unrelated servers; the
# global lock only guards the dicts.
class _StdioKeyLock:
"""A per-key lock that can be removed once nobody references it."""
def __init__(self) -> None:
self.lock = threading.Lock()
self.users = 0 # guarded by _stdio_sessions_lock
_stdio_key_locks: dict[tuple, _StdioKeyLock] = {}
_stdio_sessions_lock = threading.Lock()
_stdio_reaper_started = False
# close_stdio_sessions() can only close sessions already published in
# _stdio_sessions; one still inside connect() would be missed and cached
# stale. Bump a generation on every close so that connect discards its
# session instead of publishing it. Guarded by _stdio_sessions_lock.
_stdio_close_all_gen = 0
_stdio_url_close_gen: dict[str, int] = {}
_stdio_cfg_close_gen: dict[tuple, int] = {}
# close_stdio_sessions(url): match any env for that command.
_ANY_HEADERS = object()
def _headers_key(headers: Optional[dict]) -> tuple:
return tuple(sorted((headers or {}).items()))
def _url_close_key(url: str) -> str:
# Commands/URLs (token args, embedded credentials) and env values can hold
# secrets and these maps are never pruned; key by digest so closed/edited
# configs don't retain them in memory forever.
return hashlib.sha256(url.encode()).hexdigest()
def _cfg_close_key(url: str, headers: Optional[dict]) -> str:
return hashlib.sha256(repr((url, _headers_key(headers))).encode()).hexdigest()
def _stdio_close_generation(url: str, headers: Optional[dict]) -> tuple[int, int, int]:
return (
_stdio_close_all_gen,
_stdio_url_close_gen.get(_url_close_key(url), 0),
_stdio_cfg_close_gen.get(_cfg_close_key(url, headers), 0),
)
def _session_key(url: str, headers: Optional[dict], scope: Optional[str]) -> tuple:
return (url, _headers_key(headers), scope or "")
def _checkout_stdio_session(key: tuple) -> Optional[_StdioSession]:
session = _stdio_sessions.get(key)
if session is not None and session.is_connected():
session.last_used = time.monotonic()
session.in_flight += 1
return session
return None
def _borrow_stdio_key_lock(key: tuple) -> _StdioKeyLock:
"""Return a stable per-key lock while a caller waits for/connects it."""
key_lock = _stdio_key_locks.setdefault(key, _StdioKeyLock())
key_lock.users += 1
return key_lock
def _discard_stdio_key_lock(key: tuple) -> None:
key_lock = _stdio_key_locks.get(key)
if key_lock is not None and key_lock.users == 0 and key not in _stdio_sessions:
_stdio_key_locks.pop(key, None)
def _return_stdio_key_lock(key: tuple, key_lock: _StdioKeyLock) -> None:
with _stdio_sessions_lock:
key_lock.users -= 1
_discard_stdio_key_lock(key)
def _get_stdio_session(
url: str, headers: Optional[dict], scope: Optional[str], deadline, cancel_event, config_check
) -> _StdioSession:
"""``deadline`` is the caller's absolute monotonic budget (None = no limit):
the key-lock wait and the connect share it, so a slow startup can't stack
full timeout windows (see _call_stdio_tool)."""
global _stdio_reaper_started
key = _session_key(url, headers, scope)
with _stdio_sessions_lock:
session = _checkout_stdio_session(key)
if session is not None:
return session
key_lock = _borrow_stdio_key_lock(key)
try:
# Poll the acquire with connect()'s deadline/cancel semantics: a second
# same-scope call must not block uncancellably behind another caller's
# slow startup (e.g. a first-run npx download).
remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
# timeout=None means no key-lock deadline (only cancel unblocks it).
window = None if remaining is None else min(remaining, _STDIO_CONNECT_TIMEOUT)
lock_deadline = None if window is None else time.monotonic() + window
while not key_lock.lock.acquire(timeout = 0.05):
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
if lock_deadline is not None and time.monotonic() >= lock_deadline:
raise asyncio.TimeoutError
try:
stale = None
with _stdio_sessions_lock:
session = _checkout_stdio_session(key)
if session is not None:
return session
if key in _stdio_sessions:
stale = _stdio_sessions.pop(key)
generation = _stdio_close_generation(url, headers)
if stale is not None:
_retire_stdio_session(stale)
session = _StdioSession(url, headers)
try:
session.connect(
None if deadline is None else max(0.0, deadline - time.monotonic()),
cancel_event,
)
except Exception:
session.close()
raise
# A caller can read the server row, then lose to an update/delete whose close ran
# before our generation snapshot. Re-verify the row after connect; the generation check
# below covers a close landing between this check and publish.
if config_check is not None:
try:
current = bool(config_check())
except Exception: # noqa: BLE001
current = False
if not current:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
evicted: list = []
with _stdio_sessions_lock:
closed_while_connecting = _stdio_close_generation(url, headers) != generation
if not closed_while_connecting:
session.in_flight = 1
evicted = _evict_stdio_lru_locked() # bound the cache (LRU idle)
_stdio_sessions[key] = session
if not _stdio_reaper_started:
_stdio_reaper_started = True
threading.Thread(
target = _stdio_session_reaper, name = "mcp-stdio-reaper", daemon = True
).start()
atexit.register(close_stdio_sessions)
for victim in evicted:
logger.info("Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url))
victim.close()
if closed_while_connecting:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
return session
finally:
key_lock.lock.release()
finally:
_return_stdio_key_lock(key, key_lock)
def _release_stdio_session(session: _StdioSession) -> None:
victims: list = []
with _stdio_sessions_lock:
session.in_flight = max(0, session.in_flight - 1)
session.last_used = time.monotonic()
close_now = session.defunct and session.in_flight == 0
# Re-enforce the cap once a burst's sessions go idle. Insert-time eviction
# only trims idle sessions, so it can overshoot while every cached session
# is busy; reclaim that overshoot here instead of waiting for the idle
# reaper. Never evict the session we just used (its last_used is newest).
while len(_stdio_sessions) > _STDIO_MAX_SESSIONS:
idle = [
(s.last_used, k)
for k, s in _stdio_sessions.items()
if s.in_flight == 0 and s is not session
]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
victims.append(_stdio_sessions.pop(oldest))
_discard_stdio_key_lock(oldest)
if close_now:
session.close()
for victim in victims:
victim.close()
def _retire_stdio_session(session: _StdioSession) -> None:
"""Close a discarded session, but only once no other borrower is mid-call
on it -- overlapping same-scope calls share one client, and one call's
timeout must not kill another's in-flight request. The last borrower's
_release_stdio_session() performs the deferred close."""
with _stdio_sessions_lock:
session.defunct = True
busy = session.in_flight > 0
if not busy:
session.close()
def _drop_stdio_session(key: tuple, session: _StdioSession) -> None:
with _stdio_sessions_lock:
if _stdio_sessions.get(key) is session:
_stdio_sessions.pop(key)
_discard_stdio_key_lock(key)
_retire_stdio_session(session)
def _evict_stdio_lru_locked() -> list:
"""Caller holds _stdio_sessions_lock. Evict least-recently-used *idle*
sessions until the cache is under the cap. Returns the evicted sessions so
the caller can close them OUTSIDE the lock. If every session is busy the
cache may transiently overshoot rather than kill an in-flight call."""
victims: list = []
while len(_stdio_sessions) >= _STDIO_MAX_SESSIONS:
idle = [(s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
victims.append(_stdio_sessions.pop(oldest))
_discard_stdio_key_lock(oldest)
return victims
def close_stdio_sessions(url: Optional[str] = None, headers = _ANY_HEADERS) -> None:
"""Close persistent stdio sessions: all of them (``url`` None), every env
for one command (``headers`` omitted), or one server config (url + headers).
Two server rows can share a command with different envs; editing one must
not kill the other's live state, so the routes pass the edited row's env."""
global _stdio_close_all_gen
# HTTP/SSE servers are never cached as stdio sessions, so a specific non-stdio
# url has nothing to close and must not accrue a close-generation entry.
if url is not None and not is_stdio(url):
return
hk = None if headers is _ANY_HEADERS else _headers_key(headers)
with _stdio_sessions_lock:
if url is None:
_stdio_close_all_gen += 1
elif hk is None:
uk = _url_close_key(url)
_stdio_url_close_gen[uk] = _stdio_url_close_gen.get(uk, 0) + 1
else:
cfg = _cfg_close_key(url, headers)
_stdio_cfg_close_gen[cfg] = _stdio_cfg_close_gen.get(cfg, 0) + 1
keys = [
k
for k in _stdio_sessions
if (url is None or k[0] == url) and (hk is None or k[1] == hk)
]
sessions = [_stdio_sessions.pop(k) for k in keys]
for key in keys:
_discard_stdio_key_lock(key)
for session in sessions:
session.close()
def _reap_idle_stdio_sessions(now: Optional[float] = None) -> None:
now = time.monotonic() if now is None else now
with _stdio_sessions_lock:
expired = [
key
for key, session in _stdio_sessions.items()
if session.in_flight == 0 and now - session.last_used >= _STDIO_SESSION_IDLE_TTL
]
sessions = [_stdio_sessions.pop(key) for key in expired]
for key in expired:
_discard_stdio_key_lock(key)
for session in sessions:
logger.info("Closing idle stdio MCP session: %s", _stdio_log_id(session.url))
session.close()
def _stdio_session_reaper() -> None:
while True:
time.sleep(_STDIO_SESSION_REAP_INTERVAL)
try:
_reap_idle_stdio_sessions()
except Exception as exc: # noqa: BLE001
logger.debug("stdio session reaper iteration failed: %s", exc)
async def list_tools_async(
url: str,
headers: Optional[dict] = None,
@ -251,11 +763,10 @@ async def list_tools_async(
return await asyncio.wait_for(_fetch(), timeout = timeout)
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools()
# probes a server only on a cache miss, keeping MCP discovery off the chat
# send's critical path -- tool schemas are stable within a session. The
# /refresh route warms it; a URL/header/OAuth change or a delete evicts it.
# Successful probes are cached indefinitely.
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools() probes a server only
# on a cache miss, keeping MCP discovery off the chat send's critical path -- tool schemas are
# stable within a session. The /refresh route warms it; a URL/header/OAuth change or a delete
# evicts it. Successful probes are cached indefinitely.
_tool_cache: dict[str, list[dict]] = {}
# server_id -> monotonic time before which a failed server must not be
@ -298,23 +809,210 @@ def invalidate_tool_cache(server_id: Optional[str] = None) -> None:
_probe_cooloff_until.pop(server_id, None)
MCP_IMAGES_SENTINEL = "__MCP_IMAGES__:"
MAX_IMAGE_PAYLOAD_CHARS = 12_000_000
def _flatten_result(result: Any) -> str:
parts = []
images = []
omitted = 0
budget = MAX_IMAGE_PAYLOAD_CHARS
for block in getattr(result, "content", None) or []:
text = getattr(block, "text", None)
if text:
parts.append(str(text))
continue
data = getattr(block, "data", None)
mime = getattr(block, "mimeType", None)
if data and isinstance(mime, str) and mime.startswith("image/"):
data = str(data)
if len(data) > budget:
omitted += 1
continue
budget -= len(data)
images.append({"data": data, "mimeType": mime})
body = "\n".join(parts)
if not body:
structured = getattr(result, "structured_content", None)
body = str(structured) if structured is not None else ""
if images or omitted:
notes = []
if images:
n = len(images)
notes.append(f"{n} image{'s' if n > 1 else ''} attached; displayed to the user")
if omitted:
notes.append(f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)")
note = f"[{'; '.join(notes)}]"
body = f"{body}\n{note}" if body else note
if getattr(result, "is_error", False):
# "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge.
return f"Error: {body}" if body else "Error: tool returned no content"
body = f"Error: {body}" if body else "Error: tool returned no content"
if images:
body += "\n" + MCP_IMAGES_SENTINEL + json.dumps(images)
return body
async def _race_tool_call(call_coro, timeout: Optional[float], cancel_event) -> Any:
"""Await ``call_coro`` under ``timeout``, polling ``cancel_event`` so a
/cancel POST interrupts even mid-network-read."""
async def _watch_cancel() -> None:
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
if cancel_event is not None and cancel_event.is_set():
call_coro.close()
raise _MCPCancelled
call_task = asyncio.create_task(call_coro)
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
def _call_stdio_tool(
url: str,
headers: Optional[dict],
name: str,
args: dict,
timeout,
cancel_event,
scope: Optional[str],
config_check,
) -> Any:
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
# One deadline covers the key-lock wait, connect, call-lock wait, and the
# call itself, matching the one-shot/HTTP paths where the timeout wrapped
# connect plus call in a single window.
deadline = None if timeout is None else time.monotonic() + timeout
def _remaining() -> Optional[float]:
return None if deadline is None else max(0.0, deadline - time.monotonic())
# Callers without a Studio session id must retain the former one-shot
# behavior: no browser/cookie/tool state can leak into another request.
# Use an ephemeral key (and close it below) rather than the shared empty
# scope that the persistent-session cache used previously.
def _config_ok() -> bool:
if config_check is None:
return True
try:
return bool(config_check())
except Exception: # noqa: BLE001
return False
ephemeral = not scope
if ephemeral:
scope = f"request-{uuid.uuid4().hex}"
key = _session_key(url, headers, scope)
# attempt 0 may find the cached session stale/dead *before* dispatch and
# reconnect once (safe); attempt 1 is a freshly connected session.
for attempt in (0, 1):
session = _get_stdio_session(url, headers, scope, deadline, cancel_event, config_check)
try:
# Serialize calls per session: overlapping same-scope calls must
# not interleave operations on one stateful server (browser, REPL).
while not session.call_lock.acquire(timeout = 0.05):
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
rem = _remaining()
if rem is not None and rem <= 0:
raise asyncio.TimeoutError
except BaseException:
# Never touched the transport: keep the session for its borrower.
_release_stdio_session(session)
if ephemeral:
_drop_stdio_session(key, session)
raise
discard_session = ephemeral
retry = False
try:
# We may have waited on the call lock while another caller's timeout retired this
# session, a server update/delete invalidated it, or a reused subprocess died. Re-check
# all three before dispatch so we never run on a retired/dead client or a stale config.
if session.closed.is_set():
# Intentional close (server update/delete/shutdown): don't retry on stale config.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
elif session.defunct:
# A concurrent same-scope caller's timeout retired this session;
# move to a fresh one instead of reusing the retired client.
discard_session = True
if attempt == 0:
retry = True
else:
raise RuntimeError("MCP server session was retired during the call")
elif not _config_ok():
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
elif _transport_dead(session):
# Dead BEFORE dispatch: no request was sent, so reconnect + retry.
discard_session = True
if attempt == 0:
retry = True
else:
raise RuntimeError("MCP server connection is not available")
else:
rem = _remaining()
coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event)
return session.run(coro, rem)
except (_MCPCancelled, asyncio.TimeoutError):
# _race_tool_call cancels the pending call but cancellation is
# cooperative. Never return this client to the cache while the
# timed-out/cancelled operation might still run on its transport.
discard_session = True
raise
except _SessionWedged:
discard_session = True
raise asyncio.TimeoutError
except _SessionClosed:
# close_stdio_sessions() shut this session mid-call (server
# update/delete/shutdown); don't retry on the stale config.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
except Exception as exc:
if session.closed.is_set():
# An intentional close (server update/delete) can surface as a plain transport
# error or AttributeError instead of _SessionClosed; don't mistake it for a crash.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
# ToolError leaves the transport alive -> keep the session so its state
# survives. Any other exception is transport-level (dead subprocess,
# broken pipe): evict so it can't poison the scope, but DO NOT replay
# (the tool may already have run); the next call opens a fresh session.
if not _is_tool_error(exc):
discard_session = True
raise
finally:
# Set defunct + remove from the cache BEFORE releasing the call lock,
# so a queued same-scope borrower observes the retirement and opens a
# fresh session instead of reusing this one.
_release_stdio_session(session)
if discard_session:
_drop_stdio_session(key, session)
session.call_lock.release()
if not retry:
break
raise RuntimeError("unreachable")
def call_tool_sync(
url: str,
headers: Optional[dict],
@ -323,55 +1021,35 @@ def call_tool_sync(
timeout: Optional[float] = 300.0,
use_oauth: bool = False,
cancel_event = None,
scope: Optional[str] = None,
config_check = None,
) -> str:
"""Synchronously call an MCP tool.
"""Synchronously call an MCP tool. stdio servers reuse a persistent session
keyed by (command, env, scope) only when ``scope`` is provided; calls
without one stay one-shot. HTTP servers always stay one-shot.
``cancel_event`` (threading.Event) cancels the in-flight call when set.
``config_check`` (callable -> bool) re-validates the caller's server config
before a fresh stdio session is cached; False fails the call."""
``cancel_event``: optional ``threading.Event``. When set, the in-flight call is
cancelled and a cancellation Error returned. Polled alongside the tool call via
``asyncio.wait`` so a /cancel POST interrupts even mid-network-read.
"""
async def _call() -> Any:
async def _one_shot() -> Any:
async with _client(url, headers, use_oauth) as client:
return await client.call_tool(name, args)
async def _watch_cancel() -> None:
# 50 ms cadence keeps cancellation responsive without busy-looping;
# matches routes/inference.py's cancel watcher cadence.
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
async def _race() -> Any:
# Check cancellation before spawning the call task so a pre-set event
# short-circuits before opening the transport / HTTP connection.
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
call_task = asyncio.create_task(_call())
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
# raise_on_error=False lets an is_error result (which may still carry
# image content) reach _flatten_result instead of FastMCP raising ToolError
# and dropping the images. Transport failures still raise (handled below).
return await client.call_tool(name, args, raise_on_error = False)
try:
result = asyncio.run(_race())
if is_stdio(url):
result = _call_stdio_tool(
url, headers, name, args, timeout, cancel_event, scope, config_check
)
else:
result = asyncio.run(_race_tool_call(_one_shot(), timeout, cancel_event))
except _MCPCancelled:
return f"Error: MCP tool '{name}' cancelled"
except asyncio.TimeoutError:
return f"Error: MCP tool '{name}' timed out after {timeout:g}s"
suffix = f" after {timeout:g}s" if timeout is not None else ""
return f"Error: MCP tool '{name}' timed out{suffix}"
except Exception as exc:
logger.exception("MCP call_tool failed for %s: %s", name, exc)
return f"Error: MCP tool '{name}' failed: {exc}"

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

@ -5,14 +5,120 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm
instead of torch/transformers for model loading and generation.
"""
import json
import os
import threading
from typing import Optional, Generator
from core.inference.message_content import content_to_text
from core.inference.runtime_context import runtime_context_length
from loggers import get_logger
logger = get_logger(__name__)
def _mlx_vlm_model_config(model):
"""Return the loaded MLX model config and its type, preferring whichever of
config / _config actually carries a model_type."""
def _model_type(cfg):
return cfg.get("model_type") if isinstance(cfg, dict) else getattr(cfg, "model_type", None)
configs = [
cfg
for cfg in (getattr(model, "config", None), getattr(model, "_config", None))
if cfg is not None
]
for cfg in configs:
model_type = _model_type(cfg)
if model_type is not None:
return cfg, model_type
return (configs[0] if configs else None), None
def _render_registered_vlm_prompt(processor, model, messages, num_images):
"""Render through mlx-vlm when it declares a formatter for this model."""
from mlx_vlm import prompt_utils
config, model_type = _mlx_vlm_model_config(model)
if config is None:
return None
if model_type not in getattr(prompt_utils, "MODEL_CONFIG", {}):
return None
rendered = prompt_utils.apply_chat_template(
processor,
config,
messages,
add_generation_prompt = True,
num_images = num_images,
)
if isinstance(rendered, str) and rendered.strip():
return rendered
raise RuntimeError("mlx-vlm's registered renderer returned an empty prompt.")
def _count_vlm_images(content):
if isinstance(content, list):
return sum(_count_vlm_images(item) for item in content)
if not isinstance(content, dict):
return 0
if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"):
return 1
return _count_vlm_images(content.get("content"))
def _vlm_media_reprs(content):
if isinstance(content, list):
values = (
{str(content), json.dumps(content, ensure_ascii = False)}
if _count_vlm_images(content)
else set()
)
for item in content:
values.update(_vlm_media_reprs(item))
return values
if not isinstance(content, dict):
return set()
if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"):
return {str(content), json.dumps(content, ensure_ascii = False)}
return _vlm_media_reprs(content.get("content"))
def _prompt_serializes_vlm_media(prompt, messages):
"""Detect templates that embed the exact structured media object repr."""
media_reprs = set()
for message in messages:
if isinstance(message, dict):
media_reprs.update(_vlm_media_reprs(message.get("content")))
text_content = [
content_to_text(message.get("content")) for message in messages if isinstance(message, dict)
]
return any(
prompt.count(media_repr) > sum(content.count(media_repr) for content in text_content)
for media_repr in media_reprs
)
def _vlm_prompt_issue(prompt, messages):
if not isinstance(prompt, str) or not prompt.strip():
return "an empty prompt"
if _prompt_serializes_vlm_media(prompt, messages):
return "serialized structured image content"
return None
def _vlm_messages_have_tool_history(messages):
return any(
isinstance(message, dict)
and (
message.get("role") == "tool"
or message.get("tool_calls")
or message.get("tool_call_id")
)
for message in messages
)
def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"""Map mlx stream stats onto the usage/timings shape llama-server emits."""
prompt_n = int(prompt_n or 0)
@ -41,6 +147,48 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
}
def _mlx_distributed_rank_size(group = None):
"""Return ``(rank, world_size)`` for an optional MLX distributed group."""
if group is None:
return 0, 1
rank = int(group.rank())
world_size = int(group.size())
if world_size < 1:
raise ValueError(f"Invalid MLX distributed world_size={world_size}.")
if rank < 0 or rank >= world_size:
raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.")
return rank, world_size
def _mlx_distributed_backend_from_env():
if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"):
return "jaccl"
return None
def _init_mlx_distributed():
"""Initialize MLX distributed state, falling back to singleton metadata."""
import mlx.core as mx
group = None
rank = 0
world_size = 1
distributed = getattr(mx, "distributed", None)
init = getattr(distributed, "init", None) if distributed is not None else None
if callable(init):
backend = _mlx_distributed_backend_from_env()
if backend is None:
group = init()
else:
try:
group = init(backend = backend)
except TypeError:
group = init()
if group is not None:
rank, world_size = _mlx_distributed_rank_size(group)
return group, rank, world_size
def _make_mlx_presence_penalty_processor(penalty: float):
"""Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path.
@ -52,7 +200,7 @@ def _make_mlx_presence_penalty_processor(penalty: float):
def _processor(tokens, logits):
if state["prompt_len"] is None:
# First call = prompt only; latch its length.
# First call is prompt-only; latch its length.
state["prompt_len"] = int(tokens.shape[0])
return logits
generated = tokens[state["prompt_len"] :]
@ -61,22 +209,17 @@ def _make_mlx_presence_penalty_processor(penalty: float):
import mlx.core as mx
vocab = logits.shape[-1]
# Bound generated ids to the valid range [0, vocab) before they index
# logits. MLX does no bounds checking and out-of-bounds indexing is
# documented undefined behavior (crash / memory corruption), unlike the
# torch path's harmless negative wrap -- so this bound is load-bearing
# here and matches the torch filter seen[(seen >= 0) & (seen < vocab)].
# MLX has no boolean-mask filtering (data-dependent output shape is
# unsupported), so instead of compacting the id list we route every
# out-of-range or negative id to a scratch slot at index ``vocab`` that
# is dropped before the subtract. That scratch slot can never collide
# with a real token, so real ids (including id 0) are penalized exactly
# once and stray ids are ignored.
# Bound ids to [0, vocab) before indexing logits: MLX does no bounds
# checking and out-of-bounds indexing is undefined behavior (crash /
# corruption), unlike torch's harmless negative wrap. MLX also lacks
# boolean-mask filtering, so out-of-range/negative ids route to a
# scratch slot at index vocab (dropped before the subtract) that never
# collides with a real token: real ids (including 0) are penalized
# once, strays ignored.
valid = (generated >= 0) & (generated < vocab)
safe = mx.where(valid, generated, vocab).astype(mx.int32)
# Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate
# ids are idempotent, so presence applies once per distinct token; the
# scratch column is discarded and the full-width subtract stays on-device.
# Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are
# idempotent (presence applies once per token); scratch column dropped.
mask = mx.zeros((vocab + 1,), dtype = logits.dtype)
mask[safe] = penalty
logits = logits - mask[:vocab]
@ -93,7 +236,7 @@ class MLXInferenceBackend:
self.loaded_local_models = []
self.device = "mlx"
self._generation_lock = threading.Lock()
# usage/timings of the latest generation; shipped on gen_done.
# usage/timings of the latest generation, shipped on gen_done.
self.last_generation_stats = None
self._model = None
@ -101,6 +244,9 @@ class MLXInferenceBackend:
self._processor = None
self._is_vlm = False
self._config = {}
self._distributed_group = None
self._distributed_rank = 0
self._distributed_world_size = 1
# Recorded for unload to release pinned memory back to the OS.
self._memory_limits_applied = {}
@ -145,19 +291,26 @@ class MLXInferenceBackend:
trust_remote_code = False,
gpu_ids = None,
dtype = None,
parallel_mode = None,
distributed_group = None,
) -> bool:
import mlx.core as mx
# Keep the token so the native-template fallback can fetch a
# gated model's repo template later during generation.
# Keep the token so the native-template fallback can fetch a gated
# model's repo template during generation.
self._hf_token = hf_token
model_name = config.identifier if hasattr(config, "identifier") else str(config)
is_vision = getattr(config, "is_vision", False)
distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group)
is_distributed = distributed_group is not None and distributed_size > 1
self._distributed_group = distributed_group
self._distributed_rank = distributed_rank
self._distributed_world_size = distributed_size
# GGUF guard. GGUF models are served by llama-server in the parent
# process, not mlx-lm here. Reaching this with is_gguf=True means the
# route's first detection flaked (transient HF Hub) but the subprocess
# re-detected GGUF; raise loudly instead of a cryptic mlx_lm error.
# GGUF guard: GGUF is served by llama-server in the parent process,
# not mlx-lm. Reaching here with is_gguf=True means the route's
# detection flaked but the subprocess re-detected GGUF; raise loudly
# instead of a cryptic mlx_lm error.
if getattr(config, "is_gguf", False):
raise RuntimeError(
f"MLXInferenceBackend cannot load GGUF model '{model_name}': "
@ -176,11 +329,26 @@ class MLXInferenceBackend:
is_lora = getattr(config, "is_lora", False)
logger.info(
"Loading %s via %s (is_lora=%s)",
"Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)",
model_name,
"mlx-vlm" if is_vision else "mlx-lm",
is_lora,
is_distributed,
distributed_rank,
distributed_size,
parallel_mode,
)
if is_distributed and parallel_mode not in ("pipeline", "tensor"):
raise ValueError(
"Unsloth: distributed MLX inference requires parallel_mode='pipeline' "
"or parallel_mode='tensor'."
)
if is_distributed and is_lora:
raise ValueError(
"Unsloth: distributed MLX inference for LoRA adapter repos "
"is not supported yet. Merge/export the adapter into an MLX model "
"before distributed inference."
)
try:
from unsloth_zoo.mlx.loader import FastMLXModel
@ -190,14 +358,23 @@ class MLXInferenceBackend:
"(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon."
) from e
load_kwargs = {
"max_seq_length": max_seq_length,
"dtype": dtype,
"load_in_4bit": load_in_4bit,
"token": hf_token,
"trust_remote_code": trust_remote_code,
"text_only": False if is_vision else True,
}
if is_distributed:
if parallel_mode == "pipeline":
load_kwargs["pipeline_group"] = distributed_group
else:
load_kwargs["tensor_group"] = distributed_group
model, tokenizer_or_processor = FastMLXModel.from_pretrained(
model_name,
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
token = hf_token,
trust_remote_code = trust_remote_code,
text_only = False if is_vision else True,
**load_kwargs,
)
if is_vision:
@ -217,8 +394,7 @@ class MLXInferenceBackend:
self.models[model_name] = {
# Per-model token for the native-template fallback (matches transformers).
"hf_token": hf_token,
# Per-model consent for the native-template reload: re-use the exact
# trust_remote_code this model was loaded with (matches transformers).
# Per-model trust_remote_code reused by the native-template reload (matches transformers).
"trust_remote_code": trust_remote_code,
"model": self._model,
"tokenizer": self._tokenizer,
@ -234,8 +410,7 @@ class MLXInferenceBackend:
"has_audio_input": False,
"context_length": runtime_context_length(self._model, max_seq_length),
}
# Capture chat_template_info so the worker IPC reply ships it back and
# the route layer classifies capabilities like the other paths.
# Capture chat_template_info for the worker IPC reply and route capability classification.
self._populate_chat_template_info(model_name)
logger.info("Model %s loaded successfully", model_name)
@ -293,6 +468,9 @@ class MLXInferenceBackend:
self._model = None
self._tokenizer = None
self._processor = None
self._distributed_group = None
self._distributed_rank = 0
self._distributed_world_size = 1
if self.active_model_name == model_name:
self.active_model_name = None
gc.collect()
@ -320,8 +498,7 @@ class MLXInferenceBackend:
max_new_tokens = 256,
repetition_penalty = 1.0,
cancel_event = None,
# Reasoning / tool kwargs forwarded by the route + worker; rendered via
# apply_chat_template_for_generation like the transformers path.
# Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity).
tools = None,
enable_thinking = None,
reasoning_effort = None,
@ -334,7 +511,6 @@ class MLXInferenceBackend:
# Reset so a failed run cannot surface stale stats.
self.last_generation_stats = None
# Build messages with system prompt
full_messages = []
if system_prompt:
full_messages.append({"role": "system", "content": system_prompt})
@ -351,10 +527,7 @@ class MLXInferenceBackend:
{"type": "text", "text": content},
]
elif isinstance(content, list):
# Prepend image if not already present
has_image = any(
p.get("type") == "image" for p in content if isinstance(p, dict)
)
has_image = _count_vlm_images(content) > 0
if not has_image:
content.insert(0, {"type": "image"})
break
@ -415,6 +588,7 @@ class MLXInferenceBackend:
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
detect_think_prefill,
render_with_native_template_fallback,
)
@ -429,11 +603,11 @@ class MLXInferenceBackend:
if prompt is None:
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
# Same parity fix as the transformers backend: if the template dropped the
# requested tools, fall back to the native template so MLX text models keep
# advertising them. ``self._tokenizer`` is this entry's model_info tokenizer,
# so probe and native render share a renderer. (The VLM path renders via the
# processor for image tokens and is intentionally not wired here.)
# Parity with the transformers backend: if the template dropped the
# requested tools, fall back to the native template so MLX text models
# keep advertising them. self._tokenizer is this entry's tokenizer, so
# probe and native render share a renderer. (VLM renders via the
# processor for image tokens and is not wired here.)
model_info = self.models.get(self.active_model_name, {})
prompt = render_with_native_template_fallback(
formatted_prompt = prompt,
@ -448,6 +622,15 @@ class MLXInferenceBackend:
hf_token = model_info.get("hf_token"),
)
# An open <think> prefilled by the template lives in the prompt, not
# the generated tokens; re-emit it so the frontend renders the block.
think_prefix = detect_think_prefill(
prompt, getattr(self._tokenizer, "all_special_tokens", None)
)
# Emit it before the first token so the block renders during prefill.
if think_prefix:
yield think_prefix
sampler = make_sampler(
temp = temperature,
top_p = top_p,
@ -455,7 +638,7 @@ class MLXInferenceBackend:
min_p = float(min_p or 0.0),
min_tokens_to_keep = 1,
)
# Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths).
# Repetition and/or presence penalty processors (GGUF/safetensors parity).
logits_processors = []
if repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
@ -496,12 +679,11 @@ class MLXInferenceBackend:
):
final_response = response
token_ids.append(response.token)
# Decode full sequence with skip_special_tokens
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
)
yield cumulative
yield think_prefix + cumulative
if cancel_event and cancel_event.is_set():
break
@ -544,8 +726,7 @@ class MLXInferenceBackend:
)
# Pick the chat-template-aware caller: processors with their own
# apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it
# directly; else fall back to the nested tokenizer.
# apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer.
chat_target = self._processor
if (
getattr(self._processor, "apply_chat_template", None) is None
@ -554,28 +735,103 @@ class MLXInferenceBackend:
):
chat_target = getattr(self._processor, "tokenizer", self._processor)
prompt = apply_chat_template_for_generation(
chat_target,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# mlx_vlm's stream_generate handles pixel_values (None for text-only)
images = [image] if image is not None else None
attached_images = 0 if images is None else len(images)
structured_images = sum(
_count_vlm_images(message.get("content"))
for message in messages
if isinstance(message, dict)
)
if structured_images != attached_images:
raise RuntimeError(
f"VLM conversation contains {structured_images} structured image "
f"item(s) for {attached_images} attached image(s)."
)
prompt = None
has_tool_history = _vlm_messages_have_tool_history(messages)
prompt_error = None
try:
prompt = apply_chat_template_for_generation(
chat_target,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
if images is None or has_tool_history:
raise
prompt_error = exc
prompt_issue = (
_vlm_prompt_issue(prompt, messages) if prompt_error is None else "a rendering error"
)
if prompt_issue and has_tool_history:
raise RuntimeError(
f"VLM chat template returned {prompt_issue} and cannot be recovered "
"without dropping tool-call history."
) from prompt_error
cumulative = ""
if images is not None and prompt_issue:
if tools or any(
value is not None
for value in (enable_thinking, reasoning_effort, preserve_thinking)
):
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue} and cannot be recovered "
"without dropping requested tools or reasoning controls."
)
try:
recovered_prompt = _render_registered_vlm_prompt(
self._processor,
self._model,
messages,
len(images),
)
except Exception as recovery_error:
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue}; model-aware "
f"recovery failed: {recovery_error}"
) from recovery_error
if recovered_prompt is None:
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue}, and no registered "
"MLX VLM renderer was available for this model."
)
recovered_issue = _vlm_prompt_issue(recovered_prompt, messages)
if recovered_issue:
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"Model-aware VLM rendering returned {recovered_issue} for "
f"{attached_images} attached image(s)."
)
prompt = recovered_prompt
elif prompt_issue:
raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error
from core.inference.chat_template_helpers import detect_think_prefill
# Re-emit an open <think> prefill from the prompt (see _generate_text).
cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None))
# Emit it before the first token so the block renders during prefill.
if cumulative:
yield cumulative
logger.info(
"VLM generating: prompt_len=%d, has_image=%s",
len(prompt),
image is not None,
)
# mlx_vlm.stream_generate forwards **kwargs into generate_step, which
# builds the sampler + logits_processors internally.
# GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=``
# silently falls into **kwargs and is ignored, stuck at greedy 0.0.
# stream_generate forwards **kwargs into generate_step (builds the
# sampler + logits_processors internally). GOTCHA: generate_step expects
# temperature= (long form); temp= is silently ignored, stuck at greedy 0.0.
vlm_kwargs = dict(
max_tokens = max_new_tokens,
temperature = temperature,
@ -589,7 +845,7 @@ class MLXInferenceBackend:
)
if presence_penalty:
# Presence needs a custom processor: pass the full list (repetition +
# presence) instead of the repetition_penalty shortcut so both apply once.
# presence) instead of the repetition_penalty shortcut so both apply.
from mlx_lm.sample_utils import make_logits_processors
_vlm_processors = []
@ -634,7 +890,7 @@ class MLXInferenceBackend:
cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
# MLX LoRA adapter toggling not yet supported — generate normally
# MLX LoRA adapter toggling not yet supported; generate normally
yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs)
def reset_generation_state(self):

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

@ -50,6 +50,18 @@ _DISPATCH_DRAIN_TIMEOUT = 5.0
_UNLOAD_GEN_LOCK_TIMEOUT = 15.0
class GenStreamError(str):
"""A stream chunk carrying a real backend/generation error, not model text.
Subclasses str so existing display/logging consumers are unaffected, while
callers that must abort a distributed run on error (raise_on_streamed_error)
can distinguish a real error from model output whose visible text starts with
"Error:" by checking isinstance(chunk, GenStreamError).
"""
__slots__ = ()
class InferenceOrchestrator:
"""
Inference backend orchestrator — subprocess-based.
@ -162,6 +174,21 @@ class InferenceOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new inference subprocess."""
# Same recheck as the training/export spawns, REPAIR reservations only: a
# repair swaps without holding the lifecycle gate this load's caller owns,
# while an install cannot swap until this gate is released (and then its
# queued-load snapshot aborts it), so tolerating installs here lets the
# load win instead of failing both sides. Also covers the OpenAI
# auto-switch path, which enters _load_model_impl without route guards.
from utils.transformers_version import (
SidecarSwapInProgress,
sidecar_swap_kind,
)
if sidecar_swap_kind() == "repair":
raise SidecarSwapInProgress(
"A transformers repair is replacing the latest sidecar; retry when it completes."
)
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
@ -198,12 +225,24 @@ class InferenceOrchestrator:
if self._cancel_event is not None:
self._cancel_event.set()
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the inference subprocess."""
def is_worker_alive(self) -> bool:
"""True while the inference subprocess is running, even with no model
active (a failed load can leave a live worker holding sidecar modules)."""
proc = self._proc
return proc is not None and proc.is_alive()
def _shutdown_subprocess(self, timeout: float = 10.0) -> bool:
"""Gracefully shut down the inference subprocess.
Returns True only once the worker is confirmed dead. If it survives
terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives
SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the
pre-swap liveness guard can still observe the survivor instead of a cleared
handle and refuse the destructive sidecar swap."""
self._stop_dispatcher() # before killing subprocess
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
return True
# 1. Cancel any ongoing generation first (instant via mp.Event)
self._cancel_generation()
@ -240,12 +279,22 @@ class InferenceOrchestrator:
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
# Survived SIGKILL (uninterruptible syscall): keep the handle so callers
# and the pre-swap guard see a live worker rather than a nulled one.
logger.error(
"Inference subprocess still alive after terminate/kill; "
"preserving its handle for the pre-swap liveness check"
)
return False
self._proc = None
self._cmd_queue = None
self._resp_queue = None
self._cancel_event = None
self._drain_event = None
logger.info("Inference subprocess shut down")
return True
def _cleanup(self):
"""atexit handler."""
@ -482,13 +531,13 @@ class InferenceOrchestrator:
initial_resp_queue = self._resp_queue
while True:
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
yield f"Error: {self._subprocess_crash_message(crash_context)}"
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
return
resp = read_one(read_timeout)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield f"Error: {self._subprocess_crash_message(crash_context)}"
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
return
continue
@ -498,7 +547,7 @@ class InferenceOrchestrator:
# Subprocess-level error (no request_id); request-scoped failures
# arrive as gen_error below.
if rtype == "error" and not resp.get("request_id"):
yield f"Error: {resp.get('error', 'Unknown error')}"
yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}")
return
if rtype == "token":
@ -513,7 +562,7 @@ class InferenceOrchestrator:
stats_holder["stats"] = resp.get("stats")
return
elif rtype == "gen_error":
yield f"Error: {resp.get('error', 'Unknown error')}"
yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}")
return
# ------------------------------------------------------------------
@ -640,11 +689,11 @@ class InferenceOrchestrator:
GPU work stays serialized; this only avoids orchestrator lock contention.
"""
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess is not running"
yield GenStreamError("Error: Inference subprocess is not running")
return
if not self.active_model_name:
yield "Error: No active model"
yield GenStreamError("Error: No active model")
return
# Latch the target model so the recheck below can detect a switch that completed
# between _start_dispatcher and mailbox registration (mirrors the locked path's
@ -655,7 +704,7 @@ class InferenceOrchestrator:
# so without this early-out a compare request would enqueue a generate on the
# outgoing model and delay the switch.
if self._unload_pending:
yield "Error: model is being unloaded"
yield GenStreamError("Error: model is being unloaded")
return
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
@ -727,7 +776,7 @@ class InferenceOrchestrator:
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
if orphaned_dispatcher:
self._stop_dispatcher()
yield "Error: model is being unloaded"
yield GenStreamError("Error: model is being unloaded")
return
try:
@ -735,7 +784,7 @@ class InferenceOrchestrator:
except RuntimeError as exc:
with self._mailbox_lock:
self._mailboxes.pop(request_id, None)
yield f"Error: {exc}"
yield GenStreamError(f"Error: {exc}")
return
def read_mailbox(timeout):
@ -813,10 +862,70 @@ class InferenceOrchestrator:
self._stop_dispatcher()
return True
def share_distributed_object(
self,
obj,
timeout: Optional[float] = 300.0,
):
"""Share a small object through the worker's MLX distributed group."""
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess is not running")
self._wait_dispatcher_idle()
with self._mailbox_lock:
if self._mailboxes:
raise RuntimeError(
"Cannot share distributed objects while compare requests are active"
)
request_id = str(uuid.uuid4())
cmd = {
"type": "share_object",
"request_id": request_id,
"object": obj,
}
with self._gen_lock:
self._send_cmd(cmd)
deadline = None if timeout is None else time.monotonic() + timeout
while deadline is None or time.monotonic() < deadline:
remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout = min(remaining, 1.0))
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("sharing chat turn"))
continue
rtype = resp.get("type", "")
rid = resp.get("request_id")
if rid and rid != request_id:
logger.debug(
"Skipping response for request_id=%s while sharing request_id=%s",
rid,
request_id,
)
continue
if rtype == "shared":
return resp.get("object")
if rtype == "share_error":
raise RuntimeError(resp.get("error", "Failed to share object"))
if rtype == "error":
raise RuntimeError(resp.get("error", "Subprocess error"))
if rtype == "status":
continue
raise RuntimeError("Timeout waiting for distributed object share")
# ------------------------------------------------------------------
# Public API — same interface as InferenceBackend
# ------------------------------------------------------------------
# Monotonic count of PUBLISHED loads; lets the install route detect a load
# (including a same-model reload) that completed while it waited on the gate.
# Bumped when the load result is published, not at load start: a start-time
# bump is already visible when the installer snapshots mid-load, so the
# completed reload would look unchanged and get unloaded by the swap.
load_generation: int = 0
def load_model(
self,
config, # ModelConfig
@ -828,6 +937,8 @@ class InferenceOrchestrator:
approved_remote_code_fingerprint: Optional[str] = None,
gpu_ids: Optional[list[int]] = None,
subject: Optional[str] = None,
tensor_parallel: bool = False,
mlx_distributed: bool = False,
) -> bool:
"""Load a model for inference.
@ -853,6 +964,11 @@ class InferenceOrchestrator:
"approved_remote_code_fingerprint": approved_remote_code_fingerprint,
"subject": subject,
"gpu_ids": gpu_ids,
"tensor_parallel": bool(tensor_parallel),
"mlx_distributed": bool(mlx_distributed),
"mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline")
if mlx_distributed
else None,
}
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
gpu_ids,
@ -863,13 +979,36 @@ class InferenceOrchestrator:
sub_config["resolved_gpu_ids"] = resolved_gpu_ids
sub_config["gpu_selection"] = gpu_selection
# Recheck the sidecar reservation BEFORE tearing the old worker down,
# for REPAIRS only: an install holds this same lifecycle gate, so it
# cannot swap while this load runs, and its queued-load snapshot
# aborts it after this load publishes -- the load wins cleanly.
# Raising here (repair) keeps the current model loaded.
from utils.transformers_version import (
SidecarSwapInProgress,
sidecar_swap_kind,
)
if sidecar_swap_kind() == "repair":
raise SidecarSwapInProgress(
"A transformers repair is replacing the latest sidecar; "
"retry when it completes."
)
# Always kill the existing subprocess and spawn fresh: reusing one
# after unsloth patches torch internals breaks getsource on reload.
if self._ensure_subprocess_alive():
self._cancel_generation()
time.sleep(0.3)
self._shutdown_subprocess()
if self._shutdown_subprocess() is False:
# The worker survived terminate/kill (e.g. a wedged CUDA syscall that
# outlives SIGKILL). Its handle is kept, so is_worker_alive() and the
# pre-swap guard still see it; do not spawn a second worker over one
# still holding GPU memory. Fail so the load can retry once it exits.
raise RuntimeError(
"The current inference worker did not exit and still holds GPU "
"memory; not starting a new model over it. Retry shortly."
)
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
@ -958,6 +1097,7 @@ class InferenceOrchestrator:
return False
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
self.load_generation += 1
# A load always spawns a fresh subprocess holding only this model, so
# mirror that. A lingering stale name would pass unload_model's "not in
# self.models" guard, and the worker's absent-name fallback would unload
@ -989,8 +1129,15 @@ class InferenceOrchestrator:
self.models.clear()
raise Exception(error)
except Exception:
except Exception as exc:
self.loading_models.discard(model_name)
from utils.transformers_version import SidecarSwapInProgress
if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive():
# Raised before the old worker was torn down: the previous model
# is still live, so keep the mirrors (clearing them would let the
# installer treat the worker as inactive and kill it unreported).
raise
self.active_model_name = None
self.models.clear()
raise
@ -1221,9 +1368,11 @@ class InferenceOrchestrator:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
@ -1287,9 +1436,11 @@ class InferenceOrchestrator:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,
permission_mode = permission_mode,
)
def generate_with_adapter_control(
@ -1338,11 +1489,11 @@ class InferenceOrchestrator:
readers don't consume each other's tokens off the shared resp_queue.
"""
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess is not running"
yield GenStreamError("Error: Inference subprocess is not running")
return
if not self.active_model_name:
yield "Error: No active model"
yield GenStreamError("Error: No active model")
return
expected_model = self.active_model_name
@ -1359,7 +1510,7 @@ class InferenceOrchestrator:
# so we never generate on the wrong one.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield "Error: model is being unloaded"
yield GenStreamError("Error: model is being unloaded")
return
request_id = str(uuid.uuid4())
image_b64 = self._pil_to_base64(image) if image is not None else None
@ -1385,7 +1536,7 @@ class InferenceOrchestrator:
try:
self._send_cmd(cmd)
except RuntimeError as exc:
yield f"Error: {exc}"
yield GenStreamError(f"Error: {exc}")
return
yield from self._consume_token_stream(
@ -1544,10 +1695,10 @@ class InferenceOrchestrator:
) -> Generator[str, None, None]:
"""Shared inner logic for audio input generation (Whisper + ASR)."""
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess is not running"
yield GenStreamError("Error: Inference subprocess is not running")
return
if not self.active_model_name:
yield "Error: No active model"
yield GenStreamError("Error: No active model")
return
expected_model = self.active_model_name
@ -1556,7 +1707,7 @@ class InferenceOrchestrator:
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield "Error: model is being unloaded"
yield GenStreamError("Error: model is being unloaded")
return
request_id = str(uuid.uuid4())
@ -1583,7 +1734,7 @@ class InferenceOrchestrator:
try:
self._send_cmd(cmd)
except RuntimeError as exc:
yield f"Error: {exc}"
yield GenStreamError(f"Error: {exc}")
return
yield from self._consume_token_stream(

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

@ -41,6 +41,9 @@ _HEAL_SIGNALS = (
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
# TML Inkling native call marker (leaks as text when the server-side
# parser misses a narration-then-call turn).
"<|content_invoke_tool_json|>",
)

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

@ -15,6 +15,7 @@ parses tool calls from the cumulative text and dispatches via
"""
import bisect
import inspect
import re
import threading
from typing import Callable, Generator, Optional
@ -60,6 +61,7 @@ from core.inference.tool_loop_controller import (
status_for_tool,
tool_event_provenance,
)
from core.inference.tool_stream_exec import stream_tool_execution
from state.tool_approvals import (
TOOL_REJECTED_MESSAGE,
abort_tool_decision,
@ -402,6 +404,22 @@ def _tool_event_provenance(**flags: object) -> dict[str, object]:
return tool_event_provenance(**flags)
def _accepts_output_callback(func: Callable[..., str]) -> bool:
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
The loop's ``execute_tool`` is a parameter (tests inject fakes), so forward
the live-output kwarg only when the callable declares it or takes ``**kwargs``.
"""
try:
sig = inspect.signature(func)
except (TypeError, ValueError):
return False
params = sig.parameters
if "output_callback" in params:
return True
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
def _call_single_turn(single_turn, conversation: list, active_tools: list[dict]):
"""Call a single-turn generator with active tool schemas when supported."""
try:
@ -424,9 +442,11 @@ def run_safetensors_tool_loop(
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
@ -452,10 +472,27 @@ def run_safetensors_tool_loop(
"""
conversation = list(messages)
# Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search.
# Normalize the mode (mirrors the GGUF loop): "full" and
# bypass_permissions are the same switch; unset/unknown behaves as "ask".
# "off" keeps the sandbox but never prompts.
if permission_mode == "full":
bypass_permissions = True
elif bypass_permissions:
permission_mode = "full"
elif permission_mode not in ("ask", "auto", "off"):
permission_mode = "ask"
# Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to
# web_search. Skip only when a retrieval call would actually prompt (ask
# mode); auto never gates the safe search_knowledge_base tool.
from core.inference.tools import build_rag_autoinject
_auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope)
# off never prompts, so (like auto) it must not lose first-pass retrieval
# even if a direct caller passes a stale confirm_tool_calls flag.
_skip_autoinject = (
confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off")
)
_auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope)
if _auto:
for _ev in _auto["events"]:
yield _ev
@ -533,12 +570,24 @@ def run_safetensors_tool_loop(
provisional_render_html_started = False
provisional_resolved = False
provisional_render_html_id = f"call_{next_call_id}"
# Live-args offset for the provisional render_html card: the drained call
# text streams as tool_args so the canvas shows the HTML being written.
_live_args_streamed_upto = -1
# When a human confirmation gate is active the real tool_start is keyed
# by an approval id and carries awaiting_confirmation, so an early
# provisional card (keyed by tool_call_id, no approval) would show the
# tool as "running" before the user has approved it. Suppress the early
# card in that case and let the gated tool_start be the first signal.
_provisional_confirm_gated = bool(confirm_tool_calls) and not bypass_permissions
# In auto mode render_html is always safe and never prompts, so keep its
# early canvas card (the frontend sends confirm_tool_calls=true alongside
# auto); mirrors the GGUF path's _confirm_gated exemption.
from core.inference.tools import is_always_safe_tool
_provisional_confirm_gated = (
bool(confirm_tool_calls)
and not bypass_permissions
and not (permission_mode == "auto" and is_always_safe_tool("render_html"))
)
gen = _call_single_turn(single_turn, conversation, active_tools)
prev_cumulative = ""
@ -595,6 +644,30 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
# Backlog first: everything drained so far.
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
elif (
provisional_render_html_started
and not provisional_resolved
and _live_args_streamed_upto >= 0
and len(content_accum) > _live_args_streamed_upto
):
# Still writing the call: stream the fragment so the canvas
# renders live. Display only; content_accum still feeds the
# stream-end parser verbatim.
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum[_live_args_streamed_upto:],
}
_live_args_streamed_upto = len(content_accum)
continue
if detect_state == _state_streaming:
@ -635,6 +708,13 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
continue
cumulative_display = candidate
cleaned = strip_tool_markup_streaming(
@ -791,6 +871,13 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
continue
@ -1055,8 +1142,17 @@ def run_safetensors_tool_loop(
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
# Bypass wins over the confirm gate at the loop level too, so a
# direct internal caller passing both flags never prompts.
needs_confirm = bool(confirm_tool_calls) and not bypass_permissions
# direct internal caller passing both flags never prompts. In
# "auto" mode only calls detected as potentially unsafe pause.
# "off" never prompts (sandbox stays on).
needs_confirm = (
bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off"
)
if needs_confirm and permission_mode == "auto":
from core.inference.tools import is_potentially_unsafe_tool_call
needs_confirm = is_potentially_unsafe_tool_call(
decision.tool_name, decision.arguments
)
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
start_event = decision.tool_start_event()
@ -1109,16 +1205,30 @@ def run_safetensors_tool_loop(
):
result = RAG_SEARCH_CAP_NUDGE
else:
try:
result = execute_tool(
decision.tool_name,
decision.arguments,
# Execute in a worker thread so live stdout chunks and heartbeats
# stream while the tool blocks (the SSE route turns heartbeats into
# keepalives). execute_tool is injectable; pass output_callback
# only when it accepts it.
def _invoke_tool(_output_callback, _decision = decision):
kwargs = dict(
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
disable_sandbox = bypass_permissions,
)
if _accepts_output_callback(execute_tool):
kwargs["output_callback"] = _output_callback
return execute_tool(_decision.tool_name, _decision.arguments, **kwargs)
try:
result = yield from stream_tool_execution(
_invoke_tool,
tool_name = decision.tool_name,
tool_call_id = decision.tool_call_id,
cancel_event = cancel_event,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
result = f"Error: tool raised an exception: {exc}"

6
studio/backend/core/inference/sandbox_site/__init__.py Normal file
View file

@ -0,0 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Package marker only, so wheel builds ship this directory. It goes on the
# sandbox PYTHONPATH so site machinery imports the sibling ``sitecustomize`` at
# startup; nothing in the backend imports it directly.

313
studio/backend/core/inference/sandbox_site/sitecustomize.py Normal file
View file

@ -0,0 +1,313 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Sandbox-side compatibility shim for ChatGPT code-interpreter paths.
Models habitually write to /mnt/data (or /mnt/outputs, /home/sandbox,
/workspace), none of which exist in the Studio sandbox. This module sits on the
sandbox subprocess PYTHONPATH (see ``tools._build_safe_env``), so it loads at
interpreter startup in every sandboxed ``python`` run and any Python the
``terminal`` tool launches.
It remaps those prefixes onto the CWD in ``open`` / ``io.open``, ``os.open``,
``os.makedirs`` / ``os.mkdir`` and ``pathlib.Path.mkdir``. A write/create to a
convention prefix always heals onto the CWD; a READ heals only when the mapped
target already exists (re-reading an earlier write), so a genuinely missing
input stays truthful on the path the model used instead of silently reading a
same-basename workdir file. Since prefix lists cannot cover every invented path,
``open`` / ``io.open`` also get a create-mode fallback: an absolute path outside
the CWD whose parent is missing is redirected to the basename in the CWD. Reads
and mkdir never use the fallback (an arbitrary absolute directory can legitimately
succeed). It is collision-safe: it refuses to redirect onto an existing CWD file
(letting open raise). The patch set (io.open, os.open, os.mkdir, Path.mkdir, and
the <3.11 ``_NormalAccessor.open``) covers the low-level entry points pathlib
routes through. A one-line stderr notice fires on the first remap, and everything
is wrapped in try/except so a failure never breaks the interpreter.
Identical with and without output streaming because the child env is.
"""
import builtins
import io
import json
import os
import sys
# Code-interpreter convention prefixes. Remapping is gated on the prefix being
# ABSENT (see _remap) so a genuine host mount / user dir is never shadowed.
_PREFIXES = ("/mnt/data", "/mnt/outputs", "/home/sandbox", "/workspace")
# /tmp exists on the host; separate only to note that. The absence gate applies alike.
_CONDITIONAL_PREFIXES = ("/tmp/outputs",)
_notified = False
# Invented absolute write path -> healed CWD target, so re-writing the same
# artifact re-serves it instead of tripping the anti-clobber guard.
_remapped_writes: dict = {}
# Each tool call is a fresh subprocess (in-process map starts empty), so this
# on-disk sidecar carries the map across runs. It records only sources the
# fallback healed, so an unrelated same-basename file is never adopted.
_REMAP_SIDECAR = ".unsloth_sandbox_remap.json"
def _note(subject, original, mapped):
"""Print the one-shot stderr notice so the model learns the real location.
``subject`` is what "does not exist" (the prefix, or the whole invented
path); ``original`` is echoed in the ``(original -> mapped)`` tail.
"""
global _notified
if _notified:
return
_notified = True
print(
f"note: {subject} does not exist in this sandbox; "
f"using the working directory instead ({original} -> {mapped})",
file = sys.stderr,
)
def _contained_join(cwd, rel):
"""Join ``rel`` onto ``cwd`` so the result can never escape ``cwd``.
A habit path can carry ``..`` segments; joining verbatim would let the target
climb above the sandbox. ``..`` components are dropped and empty / ``.`` ones
ignored, keeping the result under ``cwd``.
"""
parts = []
for part in rel.split("/"):
if part == "" or part == ".":
continue
if part == "..":
if parts:
parts.pop()
continue
parts.append(part)
return os.path.join(cwd, *parts) if parts else cwd
def _map_onto_cwd(
prefix,
text,
notify = True,
):
"""Map ``<prefix>/rest`` onto ``./rest`` in the CWD, noting it once.
The suffix is contained under the CWD (see ``_contained_join``) so a path
like ``/mnt/data/../other_session/file`` cannot escape the workdir.
``notify`` is False when the caller may keep the original path (a read), so
the one-shot notice is not spent on a remap that never happens.
"""
rel = text[len(prefix) :].lstrip("/")
mapped = _contained_join(os.getcwd(), rel)
if notify:
_note(prefix, text, mapped)
return mapped
def _sidecar_path(cwd):
return os.path.join(cwd, _REMAP_SIDECAR)
def _load_sidecar(cwd):
"""Return the persisted ``source -> healed target`` map, or {} on any error
(missing/corrupt/foreign sidecar degrades to in-process-only behaviour)."""
try:
with open(_sidecar_path(cwd)) as fh:
data = json.load(fh)
except Exception: # noqa: BLE001 - a bad sidecar must never break user code
return {}
return data if isinstance(data, dict) else {}
def _record_sidecar(cwd, source, target):
"""Persist ``source -> target`` so the next run re-serves it.
Written atomically (temp + ``os.replace``) and wrapped so a read-only/full
filesystem never breaks the interpreter. The path is inside the CWD, so the
patched ``open`` leaves it untouched (no remap, no recursion).
"""
try:
data = _load_sidecar(cwd)
if data.get(source) == target:
return
data[source] = target
tmp = _sidecar_path(cwd) + ".tmp"
with open(tmp, "w") as fh:
json.dump(data, fh)
os.replace(tmp, _sidecar_path(cwd))
except Exception: # noqa: BLE001 - persistence is best effort only
pass
def _is_creating_mode(mode):
"""True only when an ``open()`` mode string can CREATE a missing file.
Only ``w`` / ``a`` / ``x`` create. ``r+`` / ``rb+`` require the path to exist,
so they must not trip the write fallback (which would corrupt an unrelated
same-basename file); ``w+`` / ``a+`` / ``x+`` still match.
"""
return isinstance(mode, str) and any(c in mode for c in ("w", "a", "x"))
def _remap_open(file, mode):
"""Remap for ``open()`` / ``io.open()``.
A prefix remap runs first: a write/create heals onto the CWD; a READ heals
only when the mapped target already exists (re-reading an earlier write),
else the original path is kept so a genuine missing input fails truthfully
instead of silently reading a same-basename workdir file. Only if no prefix
matched and the call creates does the fallback kick in: an absolute target
outside the CWD whose parent is missing is redirected to the basename in the
CWD, unless ``CWD/<basename>`` already exists (an unrelated file), in which
case the original path is kept so open raises.
"""
creating = _is_creating_mode(mode)
# notify=False: emit the notice only once we commit to the mapping below.
mapped = _remap(file, notify = False)
if mapped is not file:
# Write always heals; a read only when the mapped target exists (else keep
# the original path so a missing input stays truthful).
if creating or os.path.exists(mapped):
# Commit: emit the notice now (the notify=False peek above deferred it).
_remap(file, notify = True)
return mapped
return file
if not creating:
return file
try:
text = os.fspath(file)
except TypeError:
return file
# bytes paths left untouched (str-only, matching the prefix remaps).
if not isinstance(text, str) or not os.path.isabs(text):
return file
cwd = os.getcwd()
# Already inside the CWD: a real target the model meant; leave it alone.
if text == cwd or text.startswith(cwd + os.sep):
return file
parent = os.path.dirname(text)
# Redirect only when the parent is missing; an existing external directory is
# a deliberate target and stays truthful (os.path.exists follows symlinks).
if parent and os.path.exists(parent):
return file
base = os.path.basename(text)
# A trailing sep or '.'/'..' basename would redirect onto the CWD or its
# parent; refuse and let open raise.
if base in ("", ".", ".."):
return file
remapped = os.path.join(cwd, base)
# Never clobber an unrelated file sharing this basename (lexists catches
# dangling symlinks). But a target this fallback already healed for the same
# invented path (in-process map or cross-run sidecar) is the artifact being
# re-written, so re-serve it instead of raising on every overwrite.
if os.path.lexists(remapped) and remapped not in (
_remapped_writes.get(text),
_load_sidecar(cwd).get(text),
):
return file
_remapped_writes[text] = remapped
_record_sidecar(cwd, text, remapped)
_note(text, text, remapped)
return remapped
def _remap(path, notify = True):
"""Map ``<prefix>/rest`` onto ``./rest`` in the CWD; other paths pass through.
``notify`` is forwarded to ``_map_onto_cwd``; ``_remap_open`` passes False so
a read that keeps its original path emits no false notice.
"""
try:
text = os.fspath(path)
except TypeError:
return path
if not isinstance(text, str):
return path
for prefix in _PREFIXES + _CONDITIONAL_PREFIXES:
# Heal only while the real prefix directory is absent, so a genuine host
# mount / user directory at that prefix is never shadowed.
if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists(prefix):
return _map_onto_cwd(prefix, text, notify = notify)
return path
def _install():
import pathlib
original_open = builtins.open
original_io_open = io.open
original_os_open = os.open
original_makedirs = os.makedirs
original_mkdir = os.mkdir
original_path_mkdir = pathlib.Path.mkdir
def _open(
file,
mode = "r",
*args,
**kwargs,
):
return original_open(_remap_open(file, mode), mode, *args, **kwargs)
def _io_open(
file,
mode = "r",
*args,
**kwargs,
):
return original_io_open(_remap_open(file, mode), mode, *args, **kwargs)
# mkdir/makedirs get only the prefix remap, never the write-mode fallback:
# an arbitrary absolute directory can legitimately succeed on the host.
def _makedirs(name, *args, **kwargs):
return original_makedirs(_remap(name), *args, **kwargs)
def _mkdir(path, *args, **kwargs):
return original_mkdir(_remap(path), *args, **kwargs)
def _os_open(
path,
flags,
mode = 0o777,
*,
dir_fd = None,
):
# Path.touch() etc. go through os.open, not builtins.open. Only O_CREAT
# can create, so only it maps to "creating" mode; O_TRUNC / O_APPEND
# without O_CREAT still require the file to exist, so behave as a read.
logical_mode = "w" if (flags & os.O_CREAT) else "r"
mapped = _remap_open(path, logical_mode)
if dir_fd is None:
return original_os_open(mapped, flags, mode)
return original_os_open(mapped, flags, mode, dir_fd = dir_fd)
def _path_mkdir(self, *args, **kwargs):
# pathlib probes Path.is_dir()/os.stat (unpatched) on FileExistsError, so
# a bare os.mkdir remap would still raise when the target exists. Remap
# the receiver up front so parents/exist_ok stays idempotent.
mapped = _remap(self)
target = self if mapped is self else self.__class__(mapped)
return original_path_mkdir(target, *args, **kwargs)
builtins.open = _open
# pathlib.Path.open / write_text / read_text call io.open directly, so patch both.
io.open = _io_open
# Python < 3.11 only: pathlib's accessor captured the ORIGINAL io.open at
# import (``_NormalAccessor.open = io.open``), so the io.open patch misses it.
# Repoint it at the same wrapper (staticmethod to stay unbound); 3.11+ dropped
# the accessor, so this is a no-op there.
accessor = getattr(pathlib, "_NormalAccessor", None)
if accessor is not None and hasattr(accessor, "open"):
accessor.open = staticmethod(_io_open)
# Path.touch() and other low-level opens call os.open directly, so patch it too.
os.open = _os_open
os.makedirs = _makedirs
# Path.mkdir(parents=True) calls os.mkdir per component, so patch os.mkdir;
# patch Path.mkdir itself too so exist_ok/parents land on the mapped path.
os.mkdir = _mkdir
pathlib.Path.mkdir = _path_mkdir
try:
_install()
except Exception: # noqa: BLE001 - a broken shim must never break user code
pass

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

@ -54,6 +54,8 @@ TOOL_XML_SIGNALS = (
# Kimi K2 / Moonshot.
"<|tool_calls_section_begin|>",
"<|tool_call_begin|>",
# TML Inkling native call marker.
"<|content_invoke_tool_json|>",
)

24
studio/backend/core/inference/tool_loop_controller.py
View file

@ -233,9 +233,33 @@ def is_tool_error(result: str) -> bool:
return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
def _strip_mcp_image_suffix(result: str) -> str:
"""Drop a trailing __MCP_IMAGES__ envelope only when it is the valid JSON
image array appended by _flatten_result, so legit tool text that merely
mentions the marker is not truncated."""
head, sep, payload = result.rpartition("\n__MCP_IMAGES__:")
if not sep:
return result
try:
images = json.loads(payload)
except (ValueError, RecursionError):
return result
if not isinstance(images, list) or not images:
return result
if not all(
isinstance(img, dict)
and isinstance(img.get("data"), str)
and isinstance(img.get("mimeType"), str)
for img in images
):
return result
return head.rstrip()
def strip_result_for_model(result: str) -> str:
"""Remove frontend-only sentinels (image paths, RAG source map) before
feeding the result back to the model."""
result = _strip_mcp_image_suffix(result)
for sentinel in ("__IMAGES__:", "__RAG_SOURCES__:"):
if sentinel in result:
result = result.split(sentinel, 1)[0].rstrip()

284
studio/backend/core/inference/tool_stream_exec.py Normal file
View file

@ -0,0 +1,284 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Streaming wrapper around blocking server-side tool execution.
``stream_tool_execution`` runs a blocking tool call in a worker thread and
turns it into a generator that yields:
* ``{"type": "tool_output", "tool_name", "tool_call_id", "text"}`` -- an
incremental stdout/stderr chunk (python/terminal tools) for live UI output;
* ``{"type": "heartbeat"}`` -- emitted whenever nothing else has been yielded
for ``heartbeat_interval_s`` seconds, so the SSE route can write a
keepalive and reverse proxies (Cloudflare tunnels cap idle streams at
~100 s) never see a silent connection while a tool runs;
and *returns* the tool's final result string via ``StopIteration.value``
(``result = yield from stream_tool_execution(...)``). The returned result is
byte-identical to calling the tool directly, so tool-result parsing, nudging,
and healing downstream are untouched.
"""
from __future__ import annotations
import inspect
import queue
import threading
import time
from typing import Any, Callable, Generator
from loggers import get_logger
logger = get_logger(__name__)
def accepts_output_callback(func: Callable[..., str]) -> bool:
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
``execute_tool`` is replaceable (tests inject fakes / the pre-PR signature),
so forward the kwarg only when the callable declares it or takes ``**kwargs``
(passing it unconditionally would ``TypeError`` on an old signature).
"""
try:
params = inspect.signature(func).parameters
except (TypeError, ValueError):
return False
if "output_callback" in params:
return True
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
# Cadence of heartbeat events while a tool blocks with no output. Well under
# common proxy idle caps (Cloudflare ~100 s, nginx default 60 s).
TOOL_HEARTBEAT_INTERVAL_S = 10.0
# How often the wrapper wakes to poll for output / completion / cancellation.
_POLL_INTERVAL_S = 0.25
# Upper bound on how long teardown waits for the worker once the stream is
# closed or errors. A cancel-observing tool returns within this after
# ``cancel_event`` is set; a cancel-ignoring one is a daemon left to finish on
# its own rather than blocking teardown for the tool's full timeout.
_WORKER_JOIN_TIMEOUT_S = 5.0
# Cap on total streamed live-output characters per tool call, bounding the
# transient UI stream so a tight print loop cannot flood the SSE channel. Much
# higher than the model-visible result cap (tools._MAX_OUTPUT_CHARS) since the
# UI keeps the live stream as the displayed output when the result is truncated.
TOOL_OUTPUT_STREAM_MAX_CHARS = 400_000
_STREAM_CAPPED_NOTICE = "\n... (further live output not streamed)\n"
def _drain_queue(q: "queue.Queue", sentinel: object, max_chars: int | None) -> tuple[str, bool]:
"""Pull every currently-queued item, joining chunks in FIFO order.
With ``max_chars`` set, stop concatenating at the budget and discard the
remaining chunks in place, bounding peak allocation when a chatty tool queues
far more than the cap before the consumer wakes. The crossing chunk is sliced
to one char past the budget, enough for the caller's truncation to stay
byte-identical. Returns ``(joined_text, hit_sentinel)``; the surplus is still
scanned so completion is detected promptly.
"""
parts: list[str] = []
total = 0
dropping = False
hit_sentinel = False
while True:
try:
item = q.get_nowait()
except queue.Empty:
break
if item is sentinel:
hit_sentinel = True
break
if dropping:
continue
if max_chars is not None and total + len(item) > max_chars:
# Keep one char past the budget as the overflow signal; drop the rest.
parts.append(item[: max(0, max_chars - total) + 1])
dropping = True
continue
parts.append(item)
total += len(item)
return "".join(parts), hit_sentinel
def stream_tool_execution(
invoke: Callable[[Callable[[str], None]], str],
*,
tool_name: str,
tool_call_id: str = "",
cancel_event: Any = None,
heartbeat_interval_s: float = TOOL_HEARTBEAT_INTERVAL_S,
poll_interval_s: float = _POLL_INTERVAL_S,
) -> Generator[dict, None, str]:
"""Run ``invoke(output_callback)`` in a thread; yield live events; return the result.
``invoke`` receives a thread-safe ``callable(str)`` it may call with
incremental output chunks (or ignore entirely). Exceptions raised by the
tool propagate to the caller unchanged after the worker thread finishes.
``cancel_event`` is the request-level cancellation signal already handed to
the tool. If the consumer closes this generator early (an SSE disconnect
calls ``gen.close()``, raising ``GeneratorExit`` at a ``yield``), the wrapper
sets it so a cancel-observing tool stops, then joins the worker with a bounded
timeout. Set ONLY on that abnormal-exit path, never on a clean finish, because
the event is shared across a turn's tool calls and setting it early would
abort the next tool.
"""
output_queue: queue.Queue[Any] = queue.Queue()
done_sentinel = object()
outcome: dict[str, Any] = {}
# Bound accepted output at the PRODUCER boundary: the consumer-side cap alone
# wouldn't stop a fast worker enqueuing unboundedly while a slow SSE client
# backpressures. Accept at most one char past the cap (so the consumer still
# emits the capped notice) and drop the rest. The final result is captured
# independently, so this never changes the byte-identical result.
accepted_output_chars = 0
accepted_output_lock = threading.Lock()
def _on_output(text: str) -> None:
nonlocal accepted_output_chars
if not text:
return
with accepted_output_lock:
remaining = TOOL_OUTPUT_STREAM_MAX_CHARS + 1 - accepted_output_chars
if remaining <= 0:
return
accepted = text[:remaining]
accepted_output_chars += len(accepted)
output_queue.put(accepted)
def _run() -> None:
try:
outcome["result"] = invoke(_on_output)
except BaseException as exc: # noqa: BLE001 - re-raised on the caller side
outcome["error"] = exc
finally:
# Posted after the result/error is recorded; wakes the consumer
# immediately so fast tools pay no poll-interval latency.
output_queue.put(done_sentinel)
worker = threading.Thread(
target = _run,
daemon = True,
name = f"tool-exec-{tool_name or 'unknown'}",
)
worker.start()
# Heartbeats are paced by counting idle queue polls rather than a wall clock
# (tests patch ``time.monotonic`` globally, so the wrapper must not read it).
idle_polls_per_heartbeat = max(1, int(round(heartbeat_interval_s / poll_interval_s)))
idle_polls = 0
streamed_chars = 0
stream_capped = False
finished = False
def _drain_pending(max_chars: int | None = None) -> str:
nonlocal finished
text, hit_sentinel = _drain_queue(output_queue, done_sentinel, max_chars)
if hit_sentinel:
finished = True
return text
def _drain_and_drop() -> None:
"""Discard the current and every queued chunk without concatenating.
Past the cap every chunk is dropped, so don't pay to build a combined
string only to drop it. Still detect completion so the loop can exit.
"""
nonlocal finished
while True:
try:
item = output_queue.get_nowait()
except queue.Empty:
return
if item is done_sentinel:
finished = True
return
abnormal_exit = False
try:
while not finished:
try:
item = output_queue.get(timeout = poll_interval_s)
except queue.Empty:
# A disconnect sets cancel_event while the worker is silent;
# surface a heartbeat this poll so the route regains control and
# tears down at once, not after a full heartbeat interval.
if cancel_event is not None and cancel_event.is_set():
yield {"type": "heartbeat"}
continue
idle_polls += 1
if idle_polls >= idle_polls_per_heartbeat:
idle_polls = 0
yield {"type": "heartbeat"}
continue
if item is done_sentinel:
break
if stream_capped:
# Past the cap: drop this chunk and every queued sibling (see
# _drain_and_drop). Pace with one time.sleep per poll (not
# time.monotonic -- tests patch the clock), counted as an idle
# poll so heartbeats keep flowing while the queue stays non-empty.
_drain_and_drop()
if finished:
break
time.sleep(poll_interval_s)
idle_polls += 1
if idle_polls >= idle_polls_per_heartbeat:
idle_polls = 0
yield {"type": "heartbeat"}
continue
# Bound the join to the remaining budget so the crossing batch can't
# allocate far past the cap (surplus is truncated below anyway); the
# prefix is long enough that truncation stays byte-identical.
budget = TOOL_OUTPUT_STREAM_MAX_CHARS - streamed_chars
chunk = item + _drain_pending(max_chars = budget - len(item))
idle_polls = 0
if streamed_chars + len(chunk) > TOOL_OUTPUT_STREAM_MAX_CHARS:
chunk = chunk[: max(0, TOOL_OUTPUT_STREAM_MAX_CHARS - streamed_chars)]
chunk += _STREAM_CAPPED_NOTICE
stream_capped = True
streamed_chars += len(chunk)
if chunk:
yield {
"type": "tool_output",
"tool_name": tool_name,
"tool_call_id": tool_call_id,
"text": chunk,
}
except BaseException:
# The loop only raises when the consumer closes us early: an SSE
# disconnect calls gen.close() (GeneratorExit at the yield) or the route
# throws in. Signal cancellation so a cancel-observing tool returns; the
# daemon worker is then abandoned (see finally). Re-raise so the caller
# sees the real cause (GeneratorExit must not be swallowed). Runs ONLY on
# abnormal exit, so the shared cancel_event is never set out from under
# the next tool in a clean multi-tool turn.
abnormal_exit = True
if cancel_event is not None:
try:
cancel_event.set()
except Exception:
pass
raise
finally:
# Clean finish: the worker already recorded its result and queued the
# sentinel we consumed, so this join returns at once. Abnormal exit:
# cancel_event is set and the daemon worker abandoned, so join with a zero
# timeout -- teardown never blocks the caller (the route may close this
# generator on the event loop), and the daemon cannot outlive the process.
worker.join(timeout = 0 if abnormal_exit else _WORKER_JOIN_TIMEOUT_S)
error = outcome.get("error")
if error is not None:
raise error
# Returned verbatim (the loop's record_result handles non-str), so the
# final tool result is byte-identical to a direct execute_tool call.
return outcome.get("result")

3293
studio/backend/core/inference/tools.py
View file

File diff suppressed because it is too large Load diff

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

@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.
from __future__ import annotations
import base64
import json
from loggers import get_logger
import os
import queue as _queue
@ -26,6 +27,9 @@ from typing import Any
logger = get_logger(__name__)
from utils.hardware import apply_gpu_ids
_SHARE_OBJECT_MAX_BYTES = 1 << 20
_SHARE_OBJECT_ERROR_SIZE = -1
# studio/backend root, prepended to sys.path so the spawned subprocess can
# import the utils/core packages.
_BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent)
@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None:
logger.error("Failed to send response: %s", exc)
def _encode_share_object(obj: Any) -> bytes:
data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8")
if len(data) > _SHARE_OBJECT_MAX_BYTES:
raise ValueError("Distributed object share payload is too large")
return data
def _decode_share_object(data: Any) -> Any:
return json.loads(bytes(data.tolist()).decode("utf-8"))
def _clean_token(value: str | None) -> str | None:
"""Normalize an HF token: blank or whitespace-only becomes None."""
return value if value and value.strip() else None
@ -276,6 +291,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
hf_token = _clean_token(config.get("hf_token"))
load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True))
# Latest-transformers sidecar models load 16-bit: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (e.g. grouped-MoE torch._grouped_mm).
if load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(config["model_name"], hf_token):
load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"load (4-bit is disabled for brand-new architectures)",
config["model_name"],
)
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
trust_remote_code = True
@ -329,14 +356,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
try:
success = backend.load_model(
config = mc,
max_seq_length = config.get("max_seq_length", 2048),
load_in_4bit = load_in_4bit,
hf_token = hf_token,
trust_remote_code = trust_remote_code,
gpu_ids = config.get("resolved_gpu_ids"),
)
load_kwargs = {
"config": mc,
"max_seq_length": config.get("max_seq_length", 2048),
"load_in_4bit": load_in_4bit,
"hf_token": hf_token,
"trust_remote_code": trust_remote_code,
"gpu_ids": config.get("resolved_gpu_ids"),
}
if getattr(backend, "device", None) == "mlx":
load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode")
load_kwargs["distributed_group"] = config.get("_mlx_distributed_group")
success = backend.load_model(**load_kwargs)
finally:
heartbeat_stop.set()
@ -521,6 +552,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
)
def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None:
"""Share a small Python object across MLX distributed ranks."""
request_id = cmd.get("request_id", "")
group = getattr(backend, "_distributed_group", None)
rank = int(getattr(backend, "_distributed_rank", 0) or 0)
world_size = int(getattr(backend, "_distributed_world_size", 1) or 1)
obj = cmd.get("object")
try:
if group is None or world_size <= 1:
shared = obj
else:
import mlx.core as mx
if rank == 0:
if obj is None:
mx.eval(mx.distributed.all_sum(mx.array(0), group = group))
shared = None
else:
try:
data = mx.array(_encode_share_object(obj), dtype = mx.uint8)
except Exception:
mx.eval(
mx.distributed.all_sum(
mx.array(_SHARE_OBJECT_ERROR_SIZE),
group = group,
)
)
raise
mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group))
mx.eval(mx.distributed.all_sum(data, group = group))
shared = obj
else:
size = int(mx.distributed.all_sum(mx.array(0), group = group).item())
if size == _SHARE_OBJECT_ERROR_SIZE:
raise RuntimeError("Failed to share distributed object")
if size == 0:
shared = None
else:
data = mx.zeros(size, dtype = mx.uint8)
data = mx.distributed.all_sum(data, group = group)
shared = _decode_share_object(data)
_send_response(
resp_queue,
{
"type": "shared",
"request_id": request_id,
"object": shared,
},
)
except Exception as exc:
_send_response(
resp_queue,
{
"type": "share_error",
"request_id": request_id,
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
},
)
def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
"""Handle TTS audio generation — returns WAV bytes + sample_rate."""
request_id = cmd.get("request_id", "")
@ -720,9 +812,29 @@ def run_inference_process(
exc,
)
try:
from core.inference.mlx_inference import MLXInferenceBackend
from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed
backend = MLXInferenceBackend()
if config.get("mlx_distributed"):
group, rank, size = _init_mlx_distributed()
config["_mlx_distributed_group"] = group
if size <= 1:
# A singleton group (MLX built without distributed support,
# or an invalid launch env/hostfile) would leave nonzero ranks
# looping forever on share_distributed_object. Fail the load
# instead of silently continuing without sharding.
raise RuntimeError(
"MLX distributed launch requested but initialized a singleton "
"group (size 1). Ensure the installed MLX has distributed "
"support and the launch environment/hostfile is valid, or run "
"without distributed."
)
logger.info(
"MLX distributed initialized in worker: rank=%s size=%s mode=%s",
rank,
size,
config.get("mlx_parallel_mode"),
)
_send_response(
resp_queue,
{"type": "status", "message": "Loading model..."},
@ -764,6 +876,8 @@ def run_inference_process(
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "share_object":
_handle_share_object(backend, cmd, resp_queue)
elif cmd_type == "load":
if backend.active_model_name:
backend.unload_model(backend.active_model_name)
@ -804,7 +918,31 @@ def run_inference_process(
)
return
# ── Resolve the effective base once, before activation/gates/install (no ML import) ──
# ── Windows: check Triton availability ──
# Placed ahead of the torchao stub below (which imports torch on win32 to detect ROCm),
# matching the training and export workers' gate-then-stub ordering.
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.warning(
"Triton not found on Windows — torch.compile disabled. "
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── Stub torchao on Windows ROCm before ANY transformers import ──
# Must precede every path that pulls transformers, not just the ML imports in section 2:
# a local LoRA adapter with no recorded base reaches transformers here via
# _resolve_base_model -> utils.models. See core/_torchao_stub.py; no-op off Windows ROCm.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── Resolve the effective base once, before activation/gates/install ──
# No ML import on the common path; a local adapter with no recorded base pulls
# transformers via utils.models, which is why the stub above precedes this.
# A remote LoRA's base is in its Hub adapter_config.json (else surfaced only by ModelConfig
# after import). _lora_base is set only for a genuine adapter, never a full fine-tune's base.
import json as _json
@ -842,19 +980,7 @@ def run_inference_process(
)
return
# ── 1b. Windows: check Triton availability (must precede import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.warning(
"Triton not found on Windows — torch.compile disabled. "
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1c. Security gates, then SSM/Mamba kernels, BEFORE importing transformers ──
# ── 1b. Security gates, then SSM/Mamba kernels, BEFORE importing transformers ──
# transformers snapshots its optional-backend gates at import, so a hybrid model's kernels
# must be installed before the import below ("mamba-ssm is required" otherwise). The gates
# are metadata-only, so run them first and refuse a blocked model before any native build.
@ -977,6 +1103,9 @@ def run_inference_process(
continue
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "share_object":
_handle_share_object(backend, cmd, resp_queue)
elif cmd_type == "load":
if backend.active_model_name:
backend.unload_model(backend.active_model_name)

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

@ -21,6 +21,7 @@ from functools import lru_cache
from typing import Callable
from utils.hardware.hardware import DeviceType, get_device
from utils.transformers_dtype import dtype_kwargs
from . import config
@ -157,9 +158,7 @@ def _get(model_name: str | None = None):
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(
name, device = device, model_kwargs = {"torch_dtype": "float16"}
)
_model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
_name = name
return _model

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

@ -124,10 +124,12 @@ def apply_tool_strip_patterns(
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
# <|content_invoke_tool_json|> is TML Inkling's native call marker; its JSON uses
# an ``args`` key and the block closes with <|end_message|>.
_TC_JSON_START_RE = re.compile(r"(?:<tool_call>|<\|content_invoke_tool_json\|>)\s*\{")
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_END_TAG_RE = re.compile(r"</tool_call>|<\|end_message\|>")
_TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it.
@ -686,12 +688,24 @@ def parse_tool_calls_from_text(
if kind == "json":
obj = json.loads(content[m.end() - 1 : brace_end + 1])
name = obj.get("name", "")
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes).
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside
# Hermes) and ``args`` (TML Inkling native calls).
arguments = obj.get("arguments")
if arguments is None:
arguments = obj.get("parameters", {})
arguments = obj.get("parameters")
if arguments is None:
arguments = obj.get("args", {})
if isinstance(arguments, dict):
arguments = json.dumps(arguments)
# Inkling echoes the bare tool name (and a role opener) before the
# marker: <|message_model|>NAME<|content_invoke_tool_json|>{...}.
# Fold that echo into the markup span so promotion removes it too.
if name and content.startswith("<|content_invoke_tool_json|>", start):
pre = content[:start]
if pre.endswith(name):
start -= len(name)
if content[:start].endswith("<|message_model|>"):
start -= len("<|message_model|>")
else:
name = m.group(1)
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))

207
studio/backend/core/training/trainer.py
View file

@ -11,8 +11,10 @@ import os
import sys
import types
# Prevent tokenizer parallelism deadlocks when datasets forks.
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Off on Linux so datasets' forked map() workers can't deadlock. On spawn platforms
# (Windows/macOS) map() runs in-process, so keep the fast tokenizer's Rust threads on
# (the only parallelism single-process tokenize gets; off makes prep run serially).
os.environ["TOKENIZERS_PARALLELISM"] = "true" if sys.platform in ("win32", "darwin") else "false"
# Make compiled cache modules importable by any subprocess. On spawn platforms
# (Windows/macOS) spawned dataset.map() workers re-import top-level modules, and
@ -62,7 +64,6 @@ from loggers import get_logger
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Callable
from dataclasses import dataclass
import pandas as pd
from datasets import Dataset
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
@ -71,7 +72,7 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead
from utils.models import is_vision_model, detect_audio_type
from utils.models.model_config import _env_offline
from utils.datasets import format_and_template_dataset
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
from utils.datasets.completion_masking import apply_completion_masking
from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset
from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names
from utils.paths import (
@ -86,6 +87,11 @@ from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
from .training import (
TrainingProgress,
create_mlx_trainer_adapter,
should_use_mlx_training_backend,
)
logger = get_logger(__name__)
@ -104,31 +110,16 @@ def _build_report_targets(training_args) -> list[str] | str:
return report_to or "none"
@dataclass
class TrainingProgress:
"""Training progress tracking"""
epoch: float = 0
step: int = 0
total_steps: int = 0
loss: Optional[float] = None
learning_rate: Optional[float] = None
is_training: bool = False
is_completed: bool = False
error: Optional[str] = None
status_message: str = "Ready to train" # Current stage
elapsed_seconds: Optional[float] = None
eta_seconds: Optional[float] = None
grad_norm: Optional[float] = None
num_tokens: Optional[int] = None
eval_loss: Optional[float] = None
class UnslothTrainer:
"""
Unsloth Training Backend
"""
def __new__(cls, *args, **kwargs):
if cls is UnslothTrainer and should_use_mlx_training_backend():
return create_mlx_trainer_adapter(*args, **kwargs)
return super().__new__(cls)
def __init__(self):
self.model = None
self.tokenizer = None
@ -942,7 +933,7 @@ class UnslothTrainer:
use_gradient_checkpointing = "unsloth"
elif use_gradient_checkpointing in ("true", "1", "yes"):
use_gradient_checkpointing = True
elif use_gradient_checkpointing in ("false", "0", "no"):
elif use_gradient_checkpointing in ("false", "0", "no", "none", "off"):
use_gradient_checkpointing = False
else:
# Invalid value -> "unsloth"
@ -3466,8 +3457,6 @@ class UnslothTrainer:
# ========== TRAIN ON RESPONSES ONLY ==========
# Raw-text datasets always train on all tokens.
instruction_part = None
response_part = None
is_cpt = training_args.get("is_cpt", False)
train_on_responses_enabled = (
False
@ -3484,113 +3473,93 @@ class UnslothTrainer:
# DeepSeek OCR handles this internally in its collator, so skip
# Audio VLM handles label masking in its collator, so skip
# Markers auto-detected from the chat template first, manual table
# as fallback; gpt-oss stays on its manual markers. See
# apply_completion_masking.
if (
train_on_responses_enabled
and not self.is_audio_vlm
and not self.is_audio
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
):
try:
logger.info("Configuring train on responses only...\n")
from unsloth.chat_templates import train_on_responses_only
# Template mapping for this model
model_name_lower = self.model_name.lower()
logger.info("Configuring train on responses only...\n")
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
logger.info(f"Detected template: {template_name}\n")
def _notify(level, message):
if level == "warning":
logger.warning(message)
else:
logger.info(f"{message}\n")
if template_name in TEMPLATE_TO_RESPONSES_MAPPER:
instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][
"instruction"
]
response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"]
# No try/except: the helper handles detection failures and
# double misses itself, so an exception here is a real masking
# failure that must fail the run, not silently train on full
# sequences.
self.trainer, masking_applied = apply_completion_masking(
self.trainer,
self.model_name,
train_on_responses_only,
num_proc = config_args["dataset_num_proc"],
notify = _notify,
)
logger.info(f"Instruction marker: {instruction_part[:50]}...\n")
logger.info(f"Response marker: {response_part[:50]}...\n")
if not masking_applied:
train_on_responses_enabled = False
if masking_applied:
try:
# ── Safety net: check if all samples were filtered out ──
# train_on_responses_only masks non-response tokens with -100; a
# row becomes all -100 (Unsloth drops it) when the response
# template is not found in the formatted text. Usually a
# dataset/template mismatch (already-formatted data, or 'Train on
# completions' on data that doesn't match the model's chat
# template); only sometimes max_seq_length truncating the response
# away. Skip this len()-based check for streaming.
if detect_streaming_dataset(self.trainer.train_dataset):
logger.info("Skipping post-filter length check for streaming dataset\n")
else:
logger.info(
f"No response mapping found for template: {template_name}\n"
filtered_len = len(self.trainer.train_dataset)
original_dataset_obj = (
dataset["dataset"] if isinstance(dataset, dict) else dataset
)
train_on_responses_enabled = False
else:
logger.info(f"No template mapping found for model: {self.model_name}\n")
train_on_responses_enabled = False
except Exception as e:
logger.warning(f"Could not configure train on responses: {e}")
train_on_responses_enabled = False
# Apply train on responses only if we have valid parts
if (
train_on_responses_enabled
and instruction_part
and response_part
and not self.is_audio_vlm
and not self.is_audio
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
):
try:
from unsloth.chat_templates import train_on_responses_only
self.trainer = train_on_responses_only(
self.trainer,
instruction_part = instruction_part,
response_part = response_part,
num_proc = config_args["dataset_num_proc"],
)
logger.info("Train on responses only configured successfully\n")
# ── Safety net: check if all samples were filtered out ──
# train_on_responses_only masks non-response tokens with -100;
# a row becomes all -100 (and Unsloth drops it) when the response
# template is not found in the formatted text. That is usually a
# dataset/template mismatch (already-formatted data, or 'Train on
# completions' applied to data that doesn't match the model's chat
# template), and only sometimes max_seq_length truncating the
# response away. Skip this len()-based check for streaming.
if detect_streaming_dataset(self.trainer.train_dataset):
logger.info("Skipping post-filter length check for streaming dataset\n")
else:
filtered_len = len(self.trainer.train_dataset)
original_dataset_obj = (
dataset["dataset"] if isinstance(dataset, dict) else dataset
)
original_len = len(original_dataset_obj)
dropped = original_len - filtered_len
drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0
if filtered_len == 0 or drop_pct > 30:
max_seq = training_args.get("max_seq_length", 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) were "
f"dropped after applying 'Train on completions': after "
f"masking, those rows had no trainable response tokens "
f"left. The usual cause is that this model's response "
f"template was not found in the formatted samples, so "
f"every token was masked out. That typically means the "
f"dataset is already formatted, or its structure does "
f"not match the model's chat template, so 'Train on "
f"completions' should be turned off for this dataset. "
f"Less commonly, a max_seq_length ({max_seq}) shorter "
f"than the prompt can truncate the response away; only "
f"raise it if your samples are actually longer than that."
original_len = len(original_dataset_obj)
dropped = original_len - filtered_len
drop_pct = (
round(100 * dropped / original_len, 1) if original_len > 0 else 0
)
logger.error(error_msg)
self._update_progress(error = error_msg, is_training = False)
return
if dropped > 0:
logger.info(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
logger.info(f"Post-filter dataset size: {filtered_len} samples\n")
if filtered_len == 0 or drop_pct > 30:
max_seq = training_args.get("max_seq_length", 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) were "
f"dropped after applying 'Train on completions': after "
f"masking, those rows had no trainable response tokens "
f"left. The usual cause is that this model's response "
f"template was not found in the formatted samples, so "
f"every token was masked out. That typically means the "
f"dataset is already formatted, or its structure does "
f"not match the model's chat template, so 'Train on "
f"completions' should be turned off for this dataset. "
f"Less commonly, a max_seq_length ({max_seq}) shorter "
f"than the prompt can truncate the response away; only "
f"raise it if your samples are actually longer than that."
)
logger.error(error_msg)
self._update_progress(error = error_msg, is_training = False)
return
except Exception as e:
logger.warning(f"Failed to apply train on responses only: {e}")
train_on_responses_enabled = False
if dropped > 0:
logger.info(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
logger.info(f"Post-filter dataset size: {filtered_len} samples\n")
except Exception as e:
logger.warning(f"Post-masking dataset size check failed: {e}")
else:
if train_on_responses_enabled and is_deepseek_ocr:
logger.info("Train on responses handled by DeepSeek OCR collator\n")

1354
studio/backend/core/training/training.py
View file

File diff suppressed because it is too large Load diff

326
studio/backend/core/training/worker.py
View file

@ -1309,14 +1309,18 @@ def _normalize_mlx_studio_scheduler(value):
def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]:
"""Resolve Studio local dataset uploads without importing the GPU trainer."""
"""Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer."""
from utils.paths import resolve_dataset_path
all_files: list[str] = []
for dataset_file in file_paths or []:
file_path = (
dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file))
)
dataset_path = Path(os.path.expanduser(str(dataset_file)))
if dataset_path.is_absolute():
file_path = str(dataset_path)
elif dataset_path.exists():
file_path = str(dataset_path.resolve())
else:
file_path = str(resolve_dataset_path(str(dataset_file)))
file_path_obj = Path(file_path)
if file_path_obj.is_dir():
@ -1355,6 +1359,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str:
raise ValueError(f"Unsupported dataset format: {files[0]}")
_MLX_WORKER_COMPLETE = "_mlx_worker_complete"
def _start_mlx_stop_poller(stop_queue):
import queue as _queue
import threading
stop_save = [True]
stop_requested = [False]
trainer_ref = [None]
def is_stop_requested():
return stop_requested[0]
def poll_stop():
while True:
try:
msg = stop_queue.get(timeout = 0.25)
if msg and msg.get("type") == _MLX_WORKER_COMPLETE:
return
if msg and msg.get("type") == "stop":
stop_save[0] = msg.get("save", True)
stop_requested[0] = True
trainer = trainer_ref[0]
if trainer is not None:
trainer.stop_requested = True
return
except _queue.Empty:
continue
except (EOFError, OSError):
return
stop_thread = threading.Thread(target = poll_stop, daemon = True)
stop_thread.start()
return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread
def _resolve_mlx_output_dir(config, model_name):
from utils.paths import resolve_output_dir, default_run_dir_name
output_dir = config.get("output_dir", "")
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
return str(resolve_output_dir(output_dir))
if config.get("allow_external_output_dir"):
output_path = Path(output_dir).expanduser()
if not output_path.is_absolute():
output_path = Path.cwd() / output_path
return str(output_path.resolve())
return str(resolve_output_dir(output_dir))
def _run_mlx_training(event_queue, stop_queue, config):
"""Self-contained MLX training path for Apple Silicon.
@ -1363,8 +1419,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
"""
import time
import math
import threading
import queue as _queue
from pathlib import Path
def _send(event_type, **kwargs):
@ -1374,31 +1428,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
kwargs["message"] = sm
event_queue.put({"type": event_type, "ts": time.time(), **kwargs})
_stop_save = [True]
_stop_requested = [False]
_trainer_ref = [None]
def _is_stop_requested():
return _stop_requested[0]
def _poll_stop():
while True:
try:
msg = stop_queue.get(timeout = 1.0)
if msg and msg.get("type") == "stop":
_stop_save[0] = msg.get("save", True)
_stop_requested[0] = True
trainer = _trainer_ref[0]
if trainer is not None:
trainer.stop_requested = True
return
except _queue.Empty:
continue
except (EOFError, OSError):
return
stop_thread = threading.Thread(target = _poll_stop, daemon = True)
stop_thread.start()
_stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = (
_start_mlx_stop_poller(stop_queue)
)
_send("status", status_message = "Loading MLX libraries...")
@ -1699,6 +1731,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
# sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column).
format_type = config.get("format_type", "")
custom_format_mapping = config.get("custom_format_mapping")
dataset_final_format = ""
try:
from utils.datasets import format_and_template_dataset
def _fmt_progress(status_message = "", **_kw):
@ -1764,6 +1797,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
)
if info.get("success", True):
dataset = info.get("dataset", dataset)
dataset_final_format = str(info.get("final_format", "") or "").lower()
if eval_dataset is not None:
ev = format_and_template_dataset(
eval_dataset,
@ -1804,21 +1838,14 @@ def _run_mlx_training(event_queue, stop_queue, config):
# ── 5. Build output dir ──
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import resolve_output_dir, ensure_dir
from utils.paths import ensure_dir
output_dir = config.get("output_dir", "")
if not output_dir:
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
output_dir = _resolve_mlx_output_dir(config, model_name)
ensure_dir(Path(output_dir))
# ── 6. Create trainer ──
eval_steps_val = config.get("eval_steps", 0) or 0
if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1:
# Studio sometimes sends fraction-of-total-steps
eval_steps_val = max(1, int(eval_steps_val * max_steps))
else:
eval_steps_val = int(eval_steps_val)
@ -1869,6 +1896,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
eval_steps = eval_steps_val,
)
# Also gates the masking skip below, so defined outside the feature-detect block.
raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw"
# Feature-detect optional fields so this PR works without the paired zoo bump.
_supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {})
if "cast_norm_output_to_input_dtype" in _supported_fields:
@ -1882,7 +1912,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
if "max_grad_leaf_norm" in _supported_fields:
mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm
if "append_eos" in _supported_fields:
raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw"
# Studio SFT formatting owns rendered examples; raw/CPT text still
# needs MLX to append EOS like the CUDA raw-text path.
mlx_config_kwargs["append_eos"] = bool(raw_text_mode)
@ -1903,29 +1932,27 @@ def _run_mlx_training(event_queue, stop_queue, config):
_send("eval_configured")
# ── 7. Apply train_on_responses_only if requested ──
if config.get("train_on_completions", False):
# Auto-detect markers from the chat template first, manual table as
# fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and
# Alpaca-rendered text lacks the chat markers. Also check the resolved
# format, since format_type="auto" can land on alpaca or raw text.
if (
config.get("train_on_completions", False)
and not raw_text_mode
and format_type != "alpaca"
and dataset_final_format not in ("alpaca", "raw_text")
):
_send("status", status_message = "Configuring response-only training...")
try:
from utils.datasets import (
MODEL_TO_TEMPLATE_MAPPER,
TEMPLATE_TO_RESPONSES_MAPPER,
)
template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower())
markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None
if markers:
trainer = train_on_responses_only(
trainer,
instruction_part = markers["instruction"],
response_part = markers["response"],
)
else:
_send(
"status",
status_message = f"train_on_completions skipped (no template for {model_name})",
)
except Exception as e:
_send("status", status_message = f"train_on_completions failed: {e}")
# No catch: the helper handles detection failures and double misses, so
# an exception here is a real masking failure that must fail the run,
# not silently train on full sequences.
from utils.datasets.completion_masking import apply_completion_masking
trainer, _masking_applied = apply_completion_masking(
trainer,
model_name,
train_on_responses_only,
notify = lambda level, message: _send("status", status_message = message),
)
# ── 8. Setup wandb / tensorboard ──
wandb_run = None
@ -2043,12 +2070,27 @@ def _run_mlx_training(event_queue, stop_queue, config):
# ── 11. Run training ──
gc.collect()
mx.synchronize()
trainer.train(resume_from_checkpoint = resume_from_checkpoint)
_save_model = trainer.save_model
def _skip_internal_final_save(*args, **kwargs):
raise ValueError("worker owns final save")
trainer.save_model = _skip_internal_final_save
try:
trainer.train(resume_from_checkpoint = resume_from_checkpoint)
finally:
trainer.save_model = _save_model
# ── 12. Save and finalize ──
if trainer.stop_requested and not _stop_save[0]:
# User clicked "Cancel" (save=False) — skip saving
_send("complete", output_dir = None, status_message = "Training cancelled")
if trainer.stop_requested:
if not _stop_save[0]:
# Cancel (save=False): skip saving.
_send("complete", output_dir = None, status_message = "Training cancelled")
else:
_send("status", status_message = "Saving stopped model...")
mx.synchronize()
trainer.save_model(output_dir)
_send("complete", output_dir = output_dir, status_message = "Training stopped")
else:
_send("status", status_message = "Saving model...")
mx.synchronize()
@ -2067,6 +2109,79 @@ def _run_mlx_training(event_queue, stop_queue, config):
pass
def _is_current_process_apple_silicon() -> bool:
import platform
return platform.system() == "Darwin" and platform.machine() == "arm64"
def run_mlx_training_process(
*,
event_queue: Any,
stop_queue: Any,
config: dict,
transformers_activated: bool = False,
) -> None:
"""MLX worker entrypoint shared by Studio subprocesses and the CLI adapter."""
model_name = config["model_name"]
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.hf_xet_fallback import child_should_disable_xet
if child_should_disable_xet(config):
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
if not transformers_activated:
# Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers.
_activate_transformers_version_or_warn(model_name, config.get("hf_token") or None)
from utils.hardware import hardware as _hw
_hw.detect_hardware()
if _hw.DEVICE != _hw.DeviceType.MLX:
event_queue.put(
{
"type": "error",
"error": "MLX training requires Apple Silicon with the MLX backend available.",
"stack": "",
"ts": time.time(),
}
)
return
if config.get("is_dataset_audio"):
event_queue.put(
{
"type": "error",
"error": "Audio dataset training is not yet supported on Apple Silicon.",
"stack": "",
"ts": time.time(),
}
)
return
try:
try:
_run_mlx_training(event_queue, stop_queue, config)
finally:
try:
stop_queue.put({"type": _MLX_WORKER_COMPLETE})
except (EOFError, OSError, ValueError):
pass
except Exception as exc:
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
"""Subprocess entrypoint. Fresh Python — no stale module state.
@ -2075,7 +2190,11 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
stop_queue: mp.Queue for stop commands from the parent.
config: Training config dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Off on Linux (forked datasets map() workers deadlock otherwise); on spawn
# platforms map() is in-process, so keep tokenizer threads on for faster prep.
os.environ["TOKENIZERS_PARALLELISM"] = (
"true" if sys.platform in ("win32", "darwin") else "false"
)
os.environ["PYTHONWARNINGS"] = "ignore" # before imports
# HTTP-fallback respawn: disable Xet before any huggingface_hub import (the
@ -2141,36 +2260,26 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend
mlx_backend_requested = is_apple_silicon_training_platform()
mlx_transformers_activated = False
if mlx_backend_requested and _is_current_process_apple_silicon():
# Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers.
_activate_transformers_version_or_warn(model_name, config.get("hf_token") or None)
mlx_transformers_activated = True
from utils.hardware import hardware as _hw
_hw.detect_hardware()
if _hw.DEVICE == _hw.DeviceType.MLX:
if config.get("is_dataset_audio"):
event_queue.put(
{
"type": "error",
"error": "Audio dataset training is not yet supported on Apple Silicon.",
"stack": "",
"ts": time.time(),
}
)
return
# Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.)
# Must happen before any transformers/mlx-lm imports in _run_mlx_training.
# Non-fatal: fall through with whatever version is installed, but log
# the failure instead of swallowing it (issue #6103).
_activate_transformers_version_or_warn(model_name, config.get("hf_token") or None)
try:
_run_mlx_training(event_queue, stop_queue, config)
except Exception as exc:
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE):
run_mlx_training_process(
event_queue = event_queue,
stop_queue = stop_queue,
config = config,
transformers_activated = mlx_transformers_activated,
)
return
# ── 1. Activate correct transformers version BEFORE any ML imports ──
@ -2474,6 +2583,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
_bnb_rocm_ver,
)
# Setting BNB_ROCM_VERSION makes bitsandbytes log a benign override
# notice on import; drop only that record so real errors and mismatch
# warnings still show.
if os.environ.get("BNB_ROCM_VERSION"):
import logging as _logging
_logging.getLogger("bitsandbytes.cextension").addFilter(
lambda _r: "environment variable detected" not in _r.getMessage()
)
# Parse HIP version for the kernel-fix gate below, falling back to
# the rocm version embedded in torch.__version__ when version.hip is
# unset (AMD SDK / Radeon wheels).
@ -2693,7 +2811,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from core.training.trainer import UnslothTrainer, TrainingProgress
from core.training.training import TrainingProgress
from core.training.trainer import UnslothTrainer
from utils.paths import (
ensure_dir,
resolve_output_dir,
@ -2904,11 +3023,24 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
),
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
# Latest-sidecar models load 16-bit here too: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (same flip as the chat worker).
_train_load_in_4bit = config["load_in_4bit"]
if _train_load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(model_name, hf_token):
_train_load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"training load (4-bit is disabled for brand-new architectures)",
model_name,
)
try:
success = trainer.load_model(
model_name = model_name,
max_seq_length = config["max_seq_length"],
load_in_4bit = config["load_in_4bit"],
load_in_4bit = _train_load_in_4bit,
full_finetuning = not use_lora,
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),

4
studio/backend/hub/services/download_lifecycle.py
View file

@ -75,6 +75,10 @@ def spawn_worker(
env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
env["HF_HUB_DISABLE_TELEMETRY"] = "1"
env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
# No token in Studio settings: fall back to the backend's own HF_TOKEN so
# private repos stay downloadable (needed while inkling repos are private).
if not hf_token:
hf_token = os.environ.get("HF_TOKEN") or None
env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" if hf_token else "1"
# hf_transfer's parallel Range chunks can leave sparse partials even in
# "http" mode; disable so the worker's writer is always sequential.

65
studio/backend/hub/services/models/cache_inventory.py
View file

@ -49,6 +49,10 @@ _REPO_SIZE_NEG_TTL = 60.0
_MODEL_METADATA_TIMEOUT_SECONDS = 5.0
_repo_size_cache_lock = threading.Lock()
# Identity for a cached file with no HF blob (Windows without Developer Mode: hf
# moves the blob into snapshots/ and leaves blobs/ empty).
_LOCAL_SIZE_IDENTITY_PREFIX = "size:"
def get_repo_snapshot_metadata_cached(
repo_id: str, hf_token: Optional[str] = None
@ -135,23 +139,52 @@ def _cached_repo_file_name(file_obj) -> str:
return str(getattr(file_obj, "file_name", "")).replace("\\", "/")
def _is_real_cache_blob(blob: Optional[Path], repo_dir: Optional[Path]) -> bool:
"""True only for a real cache blob at ``<repo_dir>/blobs/<etag>``.
A no-symlink ``snapshots/`` file (name is the filename, not an etag) or a
repo's own ``blobs/`` subdir is not the cache blob store.
"""
if blob is None or repo_dir is None:
return False
try:
return blob.parent.resolve(strict = False) == (repo_dir / "blobs").resolve(strict = False)
except OSError:
return False
def _cached_blob_hash(blob_path, repo_path = None) -> Optional[str]:
"""The cache blob hash (etag) for a cached file, or None when there is no blob.
Only a real blob under the repo's ``blobs/`` dir has name == hash; a moved
no-symlink ``snapshots/`` file is "no blob", so the caller uses a size identity.
"""
path = Path(blob_path)
repo_dir = Path(repo_path) if repo_path is not None else None
return path.name if _is_real_cache_blob(path, repo_dir) else None
def local_size_identity(size: int) -> str:
"""Identity for a cached file whose blob hash is unknowable: its size.
Re-hashing multi-GB GGUFs on the inventory hot path is not viable, and a
``size:`` token never collides with a hex hash.
"""
return f"{_LOCAL_SIZE_IDENTITY_PREFIX}{int(size)}"
def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]:
"""Map each cached GGUF file's repo-relative name to the SET of its local
blob hashes across all cached revisions.
identities across all revisions.
HF names each local cache blob FILE by the file's etag (lfs.sha256 else
blob_id), so a local file's blob hash == ``Path(blob_path).name``. An updated
repo keeps BOTH the old and new revision snapshots until HF garbage-collects
them, so the same file resolves to several blobs; collecting them ALL (not
just the first one seen, since ``repo_info.revisions`` is a frozenset and
yields them in arbitrary order) lets the remote-vs-local diff treat the file
as current when the remote (``main``) blob is present in any cached revision.
Mirrors the ``cached_blob_ids`` membership test in routes/models.py.
By default this keeps the historical MAIN-GGUF-only behavior. GGUF update
checks opt into companions so a shared mmproj/MTP blob can be compared too.
An identity is the file's blob hash, or a size identity when the cache holds no
blob (Windows without Developer Mode). BOTH old and new revision blobs are kept
(a set), so the diff treats the file as current when the remote ``main`` blob is
in any cached revision. Main GGUF only by default; update checks opt into
companions to compare a shared mmproj/MTP blob too.
"""
blob_map: dict[str, set[str]] = {}
repo_path = getattr(repo_info, "repo_path", None)
for revision in repo_info.revisions:
for f in revision.files:
if include_companions:
@ -163,7 +196,13 @@ def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[
if not blob_path:
continue
name = _cached_repo_file_name(f)
blob_map.setdefault(name, set()).add(Path(blob_path).name)
identity = _cached_blob_hash(blob_path, repo_path)
if identity is None:
size = int(getattr(f, "size_on_disk", 0) or 0)
if size <= 0:
continue
identity = local_size_identity(size)
blob_map.setdefault(name, set()).add(identity)
return blob_map

8
studio/backend/hub/services/models/deletion.py
View file

@ -408,7 +408,13 @@ def reclaim_replaced_gguf_variant(
and extract_quant_label(name).lower() == variant_key,
)
for snap, blob, name in matches:
blob_hash = _blob_hash_from_path(blob) if blob is not None else None
# Prune only a file we can identify as a real, stale cache blob. A
# no-symlink snapshot file has no identifiable blob hash, so keep it.
blob_hash = (
_blob_hash_from_path(blob)
if cache_inventory._is_real_cache_blob(blob, repo_dir)
else None
)
if blob_hash is None or blob_hash in keep_main_hashes:
continue
stale_matches.append((snap, blob, name))

87
studio/backend/hub/services/models/folder_browser.py
View file

@ -15,6 +15,7 @@ from loggers import get_logger
from hub.schemas.inventory import BrowseEntry, BrowseFoldersResponse
from hub.storage.scan_folders import (
contains_sensitive_path_component,
is_denied_system_path,
list_scan_folders,
)
from hub.utils.paths import (
@ -27,7 +28,10 @@ from hub.utils.paths import (
studio_root,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from hub.services.models.common import _safe_is_dir
from hub.services.models.local_inventory import _resolve_hf_cache_dir
@ -158,8 +162,14 @@ def _looks_like_model_dir(directory: Path) -> bool:
return False
def _build_browse_allowlist() -> list[Path]:
"""Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary."""
def _build_browse_allowlist(
media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None
) -> list[Path]:
"""Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary.
*media_roots* / *drive_roots* let the caller pass already-probed
removable-media and Windows drive roots so they aren't scanned again (a
disconnected mapped drive can make each probe slow); probed here when ``None``."""
from hub.storage.scan_folders import list_scan_folders
candidates: list[Path] = []
@ -176,7 +186,13 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
if media_roots is None:
media_roots = linux_run_media_mount_roots()
if drive_roots is None:
drive_roots = windows_drive_roots()
for p in media_roots:
_add(p)
for p in drive_roots:
_add(p)
_add(_resolve_hf_cache_dir())
try:
@ -218,7 +234,14 @@ def _build_browse_allowlist() -> list[Path]:
def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
"""True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox."""
"""True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox.
A Windows drive root (``D:\\``) authorizes its descendants, but a bare POSIX
root (``/``) must NOT: a single ``/`` allowlist entry (e.g. a legacy scan
folder) would otherwise authorize every absolute path, reaching ``/var``,
``/root``, etc. the denylist does not cover. Mirrors the legacy browser so
both treat ``/`` identically.
"""
try:
target_real = os.path.normcase(os.path.realpath(str(target)))
except OSError:
@ -228,13 +251,25 @@ def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
root_real = os.path.normcase(os.path.realpath(str(root)))
except OSError:
continue
if target_real == root_real:
return True
drive, tail = os.path.splitdrive(root_real)
if os.path.dirname(root_real) == root_real and not drive:
# Bare POSIX filesystem root ("/"): equality above is the only
# match; do not let it authorize arbitrary descendants.
continue
if drive.startswith(("\\\\", "//")) and not tail:
# Bare UNC share root (\\server\share): os.path.commonpath raises
# "can't mix absolute and relative" on it, so authorize its
# descendants with a boundary-safe prefix test (normcase applied).
if target_real.startswith(root_real.rstrip("\\/") + os.sep):
return True
continue
try:
if os.path.commonpath([target_real, root_real]) == root_real:
return True
except ValueError:
continue
if target_real == root_real:
return True
return False
@ -347,6 +382,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if is_denied_system_path(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
@ -354,6 +394,13 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
# Zero-component case: the requested path IS an allowlist root
# (e.g. a legacy-registered "/" or a Windows drive root).
if is_denied_system_path(str(current)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -382,9 +429,13 @@ def browse_folders_response(
"""
from hub.storage.scan_folders import list_scan_folders
# Probe removable-media and Windows drive roots once; the allowlist and
# chips reuse the result so a disconnected mapped drive isn't scanned twice.
media_roots = linux_run_media_mount_roots()
drive_roots = windows_drive_roots()
# Build the allowlist once -- the sandbox check and suggestion chips share
# it so chips are always navigable.
allowed_roots = _build_browse_allowlist()
allowed_roots = _build_browse_allowlist(media_roots, drive_roots)
try:
target = _resolve_browse_target(path, allowed_roots)
@ -440,6 +491,15 @@ def browse_folders_response(
# descending into them is refused and registration rejects them.
if contains_sensitive_path_component(name):
continue
# Same for denied system dirs (C:\Windows, /etc, ...): descent 403s,
# so don't render them as clickable rows. Resolve first so a
# symlink/junction into a denied dir is hidden too, not just a literal name.
try:
resolved_child = os.path.realpath(str(child))
except (OSError, ValueError):
resolved_child = str(child)
if is_denied_system_path(resolved_child):
continue
entries.append(
BrowseEntry(
name = name,
@ -487,13 +547,22 @@ def browse_folders_response(
return
if resolved in seen_sug:
return
# Drop a denied system dir (e.g. a stale scan-folder row) so it never
# becomes a chip that 403s on click. Drive roots stay: only their
# system subdirectories are denied, not the root itself.
if is_denied_system_path(resolved):
return
if _safe_is_dir(resolved):
seen_sug.add(resolved)
suggestions.append(resolved)
# Home first as the safe fallback.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
# Reuse the roots probed for the allowlist above (no second drive scan).
for p in media_roots:
_add_sug(p)
# Windows drive roots so the user can hop between C:, D:, E: ...
for p in drive_roots:
_add_sug(p)
# The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default.
try:

23
studio/backend/hub/services/models/gguf_variants.py
View file

@ -337,6 +337,22 @@ def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str
return result
def _size_identity_matches(local_set: set[str], remote_size: int) -> bool:
"""Whether a cached file with NO blob hash is current, judged by size.
A size token only lands in ``local_set`` for a file the cache has no blob for,
so it never loosens the hash comparison for a normal file. Tradeoff: an
equal-size requant is missed, versus the status quo where every no-blob GGUF
shows a phantom update that no re-download clears.
"""
size = int(remote_size or 0)
if size <= 0:
return False
from hub.services.models import cache_inventory
return cache_inventory.local_size_identity(size) in local_set
def _variant_update_available_from_requirement(
local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str
) -> bool:
@ -355,8 +371,13 @@ def _variant_update_available_from_requirement(
if not remote_blob:
continue
local_set = local_by_posix.get(path)
if not local_set or remote_blob not in local_set:
if not local_set:
return True
if remote_blob in local_set:
continue
if _size_identity_matches(local_set, expected.size):
continue
return True
return False

25
studio/backend/hub/services/snapshot_progress.py
View file

@ -12,6 +12,7 @@ summing stale blobs against the wrong total)."""
from __future__ import annotations
import asyncio
import threading
from pathlib import Path
from typing import Callable, Optional
@ -34,6 +35,28 @@ logger = get_logger(__name__)
# (repo_id, hf_token) -> (expected_total_bytes, expected_blob_hashes)
SnapshotMetadataResolver = Callable[[str, Optional[str]], "tuple[int, frozenset[str]]"]
# One progress log per 10% step per job, so an active download reports progress
# without emitting a line on every poll.
_progress_step_lock = threading.Lock()
_last_progress_step: dict[str, int] = {}
def _log_progress_step(job_key: str, repo_id: str, variant: Optional[str], progress: float) -> None:
step = int(progress * 10)
with _progress_step_lock:
last = _last_progress_step.get(job_key, -1)
if step == last:
return
_last_progress_step[job_key] = step
if step < last:
return # download restarted; resync without logging
logger.info(
"hub_download_progress",
repo_id = repo_id,
variant = variant or "",
percent = step * 10,
)
def _empty_progress(expected_bytes: int) -> dict:
return {
@ -215,6 +238,8 @@ def compute_snapshot_progress(
else 0
)
)
if force_active:
_log_progress_step(job_key, repo_id, variant, progress)
return {
"downloaded_bytes": display_downloaded_bytes,
"completed_bytes": display_completed_bytes,

26
studio/backend/hub/storage/scan_folders.py
View file

@ -16,7 +16,7 @@ from datetime import datetime, timezone
from storage.studio_db import get_connection
from hub.utils.paths import normalize_path
from utils.paths.external_media import is_linux_run_media_path
from utils.paths.external_media import is_linux_run_media_path, is_local_filesystem_root
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
@ -52,6 +52,25 @@ def _denied_path_prefixes() -> list[str]:
return []
def is_denied_system_path(path: str) -> bool:
"""True if *path* is, or descends from, a denied system directory.
Mirrors the denylist add_scan_folder() enforces at registration so the
browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds
a broad root (a Windows drive root C:\\ or a legacy-registered / root). The
/run carve-out keeps Linux removable-media mounts browseable. Expects an
already-resolved (realpath) path so symlinks cannot escape into a denied subtree.
"""
is_win = platform.system() == "Windows"
check = os.path.normcase(path) if is_win else path
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
if prefix == "/run" and is_linux_run_media_path(check):
continue
return True
return False
def _contains_sensitive_path_component(path: str) -> bool:
return _shared_contains_sensitive_path_component(path)
@ -108,8 +127,9 @@ def add_scan_folder(path: str) -> dict:
raise ValueError("Path must be a directory, not a file")
if not os.access(normalized, os.R_OK | os.X_OK):
raise ValueError("Path is not readable")
if os.path.dirname(normalized) == normalized:
# Registering a filesystem root would expose denied system dirs via browse.
if is_local_filesystem_root(normalized):
# A local fs root ("/", "C:\\") would expose denied system dirs via browse;
# a UNC share root (\\server\share) has none under it and stays registerable.
raise ValueError("The filesystem root cannot be registered")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")

63
studio/backend/hub/tests/test_model_services.py
View file

@ -36,6 +36,20 @@ from hub.utils import (
from hub.workers import hf_download
@pytest.fixture(autouse = True)
def _denylist_inert(monkeypatch):
# The browse tests here exercise allowlist containment, symlink safety and
# the sensitive-name filter, not the system-directory denylist (which has
# its own suite in tests/test_browse_denylist.py). On macOS tmp_path
# resolves under /private/var, a denied prefix, so _resolve_browse_target
# would 403 the fixture dirs before that logic runs. Keep the denylist inert
# so these assertions hold on every platform. folder_browser binds
# is_denied_system_path at import, so patch it on that module, not on
# scan_folders. The "rejects" cases still 403 via the allowlist/sensitive
# checks, and the non-browse tests never call it.
monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False)
def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path):
return SimpleNamespace(
repo_id = repo_id,
@ -105,6 +119,10 @@ def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id):
assert paths.is_valid_repo_id(repo_id)
def test_repo_id_validation_accepts_max_length_namespaced_repo():
assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}")
@pytest.mark.parametrize(
"repo_id",
[
@ -121,6 +139,48 @@ def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id):
assert not paths.is_valid_repo_id(repo_id)
def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M")
assert path is not None
assert path.name == "models--owner--repo--variant--q4_k_m.json"
@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64])
def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
repo_id = f"{'a' * 96}/{'b' * 96}"
assert paths.is_valid_repo_id(repo_id)
assert download_manifest.write_cancel_marker("model", repo_id, variant, "http")
assert download_manifest.write_manifest(
"model",
repo_id,
variant,
[download_manifest.ExpectedFile(path = "model.gguf", size = 1)],
"http",
)
marker_path = state_dir.marker_path("model", repo_id, variant)
manifest_path = state_dir.manifest_path("model", repo_id, variant)
assert marker_path is not None
assert manifest_path is not None
assert "--sha256-" in marker_path.name
assert len(marker_path.name.encode("utf-8")) <= 255
assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255
assert download_manifest.has_cancel_marker("model", repo_id, variant)
assert download_manifest.read_manifest("model", repo_id, variant) is not None
assert list(download_manifest.iter_variant_markers("model", repo_id)) == [
(variant, marker_path)
]
assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [
(variant, manifest_path)
]
class _RecordingLogger:
def __init__(self):
self.warnings = []
@ -182,7 +242,8 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
home = tmp_path / "home"
(home / ".ssh").mkdir(parents = True)
(home / "models").mkdir()
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda: [home])
# Accept and ignore the optional (media_roots, drive_roots) args the caller now passes.
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home])
response = folder_browser.browse_folders_response(str(home), show_hidden = True)

13
studio/backend/hub/utils/gguf_plan.py
View file

@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
"""The separate MTP drafter to fetch with every variant: the repo-root
``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
(Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
prefix, first in sort order) so download and load resolve the same file;
the higher-precision ``MTP/`` subdir copies are for explicit selection and
are not auto-fetched. None for repos with the head baked into the main
GGUF (Qwen)."""
(Gemma 4). Same pick as the loader's drafter resolution (root-level
``mtp-`` prefix, first in sort order) so download and load resolve the same
file; the higher-precision ``MTP/`` subdir copies are for explicit
selection and are not auto-fetched. None for repos with the head baked into
the main GGUF (Qwen)."""
# Root-level only: the MTP/ subdir copies now share the mtp- prefix too.
candidates = sorted(
(
s
for s in siblings
if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-")
),
key = lambda s: getattr(s, "rfilename"),
)

9
studio/backend/hub/utils/paths.py
View file

@ -181,15 +181,20 @@ def is_valid_repo_id(repo_id: str) -> bool:
"""Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs."""
if not repo_id or repo_id != repo_id.strip():
return False
if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"):
if repo_id.endswith(".git"):
return False
if "--" in repo_id or ".." in repo_id:
return False
segments = repo_id.split("/")
if len(segments) not in (1, 2):
return False
# Match huggingface_hub.validate_repo_id: the 96-char limit applies per
# segment (repo name / namespace), not to the whole "namespace/repo_name"
# string, so long-but-valid repo names are not falsely rejected.
return all(
segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None
segment not in ("", ".", "..")
and len(segment) <= _MAX_REPO_ID_LENGTH
and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None
for segment in segments
)

33
studio/backend/hub/utils/state_dir.py
View file

@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories:
manifests/ <key>.json per-download expected-files manifest
cancelled/ <key>.json per-download cancel marker
The ``<key>`` mirrors HF's cache dir naming so a state file can be
eyeballed next to the on-disk repo it describes:
The ``<key>`` mirrors HF's cache dir naming while the resulting manifest,
cancel-marker, and atomic-write temp filenames fit common filesystem basename
limits. Very long repo IDs use a stable hash in the state key:
models--<owner>--<name> full snapshot
models--<owner>--<name>--variant--<variant> GGUF variant
@ -49,6 +50,11 @@ _MANIFESTS_SUBDIR = "manifests"
_CANCELLED_SUBDIR = "cancelled"
_WORKERS_SUBDIR = "workers"
_SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$")
_MAX_STATE_BASENAME_BYTES = 255
_STATE_EXTENSION = ".json"
# _atomic_write_json writes ".<target>.tmp-<8hex>" beside the final file.
_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8
_MAX_VARIANT_FRAGMENT_LENGTH = 64
def state_root() -> Optional[Path]:
@ -84,16 +90,35 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str:
return f"{repo_type}s--{repo_id.replace('/', '--')}".lower()
def _filename_bytes(name: str) -> int:
return len(name.encode("utf-8"))
def _state_filename_fits(entry_key: str) -> bool:
filename = f"{entry_key}{_STATE_EXTENSION}"
return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES
def _state_repo_key(repo_type: RepoType, repo_id: str) -> str:
base = repo_cache_basename(repo_type, repo_id)
variant_prefix = f"{base}--variant--"
longest_variant_key = f"{variant_prefix}{'x' * _MAX_VARIANT_FRAGMENT_LENGTH}"
if _state_filename_fits(longest_variant_key):
return base
digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:32]
return f"{repo_type}s--sha256-{digest}"
def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str:
"""Lowercased prefix every variant-keyed state file for this repo shares.
The single source the download_manifest enumerators match against, so the
scheme in :func:`_entry_key` cannot drift from them silently."""
return f"{repo_cache_basename(repo_type, repo_id)}--variant--"
return f"{_state_repo_key(repo_type, repo_id)}--variant--"
def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str:
base = repo_cache_basename(repo_type, repo_id)
base = _state_repo_key(repo_type, repo_id)
if not variant:
return base
normalized_variant = variant.strip().lower()

18
studio/backend/loggers/config.py
View file

@ -17,6 +17,15 @@ import structlog
from loggers.handlers import filter_sensitive_data
class _DropTorchDtypeDeprecation(logging.Filter):
"""Drop transformers' once-per-run "`torch_dtype` is deprecated" warning_once.
It is emitted via logging (not warnings), so a warnings filter cannot catch it."""
def filter(self, record: logging.LogRecord) -> bool:
msg = record.getMessage()
return not ("torch_dtype" in msg and "deprecated" in msg)
class LogConfig:
"""Structured logging configuration for the application."""
@ -72,4 +81,13 @@ class LogConfig:
cache_logger_on_first_use = True,
)
# Drop transformers' cosmetic "`torch_dtype` is deprecated" warning_once (see filter).
_dtype_filter = _DropTorchDtypeDeprecation()
for _name in (
"transformers.configuration_utils",
"transformers.modeling_utils",
"transformers.pipelines.base",
):
logging.getLogger(_name).addFilter(_dtype_filter)
return structlog.get_logger(service_name)

81
studio/backend/loggers/handlers.py
View file

@ -28,19 +28,26 @@ def _env_int(name: str, default: int) -> int:
return default
# Drop duplicate successful-GET access logs repeated within the window: the SPA
# fans one cache invalidation into many identical list fetches; only the first
# informs. Loading polls, mutations, and errors are unaffected. 0 = log all.
# Collapse identical GET/2xx logs within the window (the SPA fans one invalidation
# into many list fetches). Mutations and errors always log. 0 = off.
_ACCESS_LOG_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS", 300)
# Pure-liveness/UI polls whose access line carries no signal beyond "client still
# polling" (state changes are logged by their own modules). Collapsed to a longer
# heartbeat instead of one line per poll; first hit and any error still log. 0 = off.
# Liveness/UI polls whose line means only "still polling"; collapse to a longer
# heartbeat. First hit and errors still log. 0 = off.
_QUIET_POLL_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS", 10000)
_QUIET_POLL_PATHS = {
"/api/health",
"/api/auth/status",
"/api/inference/status",
"/api/inference/monitor",
# List polls the tabs refetch on a timer and on every tab switch.
"/api/train/runs",
"/api/models/checkpoints",
"/api/models/local",
"/api/rag/knowledge-bases",
# Legacy download polls emit no progress events (unlike /api/hub/*), so heartbeat them.
"/api/models/download-progress",
"/api/models/gguf-download-progress",
"/api/datasets/download-progress",
}
_DEDUP_MAP_MAX = 4096
_NATIVE_PATH_LEASE_RE = re.compile(
@ -62,6 +69,46 @@ _EXCLUDED_SUFFIXES = (
".woff2",
".ttf",
)
# GET polls whose 2xx line carries no signal (their progress/phase events and the UI
# do), so drop it entirely; non-2xx still logs. Only /api/hub download polls emit
# events; the legacy /api/models and /api/datasets ones heartbeat via _QUIET_POLL_PATHS.
_QUIET_SUCCESS_PATHS = {
"/api/inference/load-progress",
"/api/llama/update-status",
"/api/export/logs",
"/api/export/status",
"/api/hub/download-status",
"/api/hub/download-progress",
"/api/hub/gguf-download-progress",
"/api/hub/active-downloads",
"/api/hub/transport-status",
"/api/hub/datasets/download-status",
"/api/hub/datasets/download-progress",
"/api/hub/datasets/active-downloads",
"/api/hub/datasets/transport-status",
}
# The token-refresh route. Its first 2xx means the client has obtained a valid
# session, so from then on chat 401s are real failures and must stay visible.
_AUTH_REFRESH_PATH = "/api/auth/refresh"
# High-frequency chat list polls; their 2xx is covered by generation/tool-call/stats
# events. Exact paths only, so detail/message reads (/threads/{id}, .../messages,
# /projects/{id}) keep their logs. The pre-auth 401 race also fires on these polls.
_CHAT_LIST_PATHS = {
"/api/chat/threads",
"/api/chat/projects",
}
def _is_quiet_success(method: str, path: str, status_code: int, pre_auth: bool) -> bool:
"""GET-only. Suppress a 2xx poll line that carries no signal, plus a chat list
poll's transient pre-auth 401 (only in the bootstrap window before the first
successful token refresh). Mutations, real (post-refresh) auth failures, and
all other errors always log."""
if method != "GET":
return False
if 200 <= status_code < 300:
return path in _QUIET_SUCCESS_PATHS or path in _CHAT_LIST_PATHS
return pre_auth and status_code == 401 and path in _CHAT_LIST_PATHS
class LoggingMiddleware:
@ -71,14 +118,16 @@ class LoggingMiddleware:
self.app = app
# (method, path, query, status_code) -> monotonic ts of the last EMITTED log.
self._last_log: dict[tuple[str, str, bytes, int], float] = {}
# Flips True after the first successful /api/auth/refresh; before that, chat
# list-poll 401s are the transient bootstrap race and are suppressed.
self._auth_refreshed = False
def _is_redundant_repeat(
self, method: str, path: str, query: bytes, status_code: int, now: float
) -> bool:
"""True if an identical GET/2xx log fired < window ago. The query string
is part of the identity, so distinct query-driven GETs are not collapsed.
Mutations and non-2xx are never deduped. Quiet-poll paths use a longer
heartbeat window. Stamps only on emit, so steady polls still log."""
"""True if an identical GET/2xx log fired < window ago (query string is part
of the identity). Non-GET/non-2xx never dedup; quiet-poll paths use the longer
heartbeat. Stamps only on emit, so steady polls still log."""
if method != "GET" or not (200 <= status_code < 300):
return False
window_ms = _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS
@ -129,8 +178,16 @@ class LoggingMiddleware:
raise
else:
end_time = time.perf_counter()
if not excluded and not self._is_redundant_repeat(
scope["method"], path, scope.get("query_string", b""), status_code, end_time
if 200 <= status_code < 300 and path == _AUTH_REFRESH_PATH:
self._auth_refreshed = True
if (
not excluded
and not _is_quiet_success(
scope["method"], path, status_code, not self._auth_refreshed
)
and not self._is_redundant_repeat(
scope["method"], path, scope.get("query_string", b""), status_code, end_time
)
):
logger.info(
"request_completed",

91
studio/backend/main.py
View file

@ -159,6 +159,14 @@ if sys.platform == "win32":
_bnb_rocm_ver_final,
)
# Setting BNB_ROCM_VERSION makes bitsandbytes log a benign override notice on
# import; drop only that record so real errors and mismatch warnings show.
if os.environ.get("BNB_ROCM_VERSION"):
import logging as _logging
_logging.getLogger("bitsandbytes.cextension").addFilter(
lambda _r: "environment variable detected" not in _r.getMessage()
)
# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ──────
# In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over
# /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE
@ -226,6 +234,7 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
import hashlib
import ipaddress
import mimetypes
import re as _re
import shutil
@ -551,8 +560,12 @@ async def lifespan(app: FastAPI):
(_time.perf_counter() - _lifespan_started) * 1000,
)
# run_server's pre-bind gate sets suppress_bootstrap_injection when a public
# URL is about to serve with the default credential active: never (re)capture
# the bootstrap password into app.state, or the HTML would hand it out.
_suppress_bootstrap = getattr(app.state, "suppress_bootstrap_injection", False)
if storage.ensure_default_admin():
bootstrap_pw = storage.get_bootstrap_password()
bootstrap_pw = None if _suppress_bootstrap else storage.get_bootstrap_password()
app.state.bootstrap_password = bootstrap_pw
bootstrap_path = storage.DB_PATH.parent / ".bootstrap_password"
@ -563,7 +576,9 @@ async def lifespan(app: FastAPI):
print(" Open the Studio UI to sign in and change it.")
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
app.state.bootstrap_password = (
None if _suppress_bootstrap else storage.get_bootstrap_password()
)
_lifespan_log.info(
"lifespan startup completed in %.1fms",
@ -1355,6 +1370,61 @@ def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]
return (scheme, host, port)
def _is_loopback_ip(host: Optional[str]) -> bool:
"""Return whether ``host`` is a loopback IP, including IPv4-mapped IPv6."""
if not host or "%" in host: # a scope id (::1%eth0) is never a plain loopback
return False
try:
ip = ipaddress.ip_address(host)
except (TypeError, ValueError):
return False
mapped = getattr(ip, "ipv4_mapped", None)
return ip.is_loopback or (mapped is not None and mapped.is_loopback)
# A loopback peer carrying any of these is a proxy/tunnel relaying a remote
# client, so the peer is the proxy, not the caller: cloudflared sets
# cf-connecting-ip, reverse proxies set the rest (uvicorn only consumes
# x-forwarded-for, so the others survive to here).
_PROXIED_CLIENT_HEADERS = (
"cf-connecting-ip",
"forwarded",
"x-forwarded-for",
"x-forwarded-host",
"x-real-ip",
)
def _host_header_is_loopback(host_header: Optional[str]) -> bool:
"""Loopback/localhost check on the raw Host header.
Reads the header directly so a malformed or absent Host cannot fall back to
``request.url.hostname``'s (loopback) ASGI server address.
"""
if not host_header:
return False
host = host_header.strip()
if host.startswith("["): # [IPv6] or [IPv6]:port
end = host.find("]")
if end == -1 or (host[end + 1 :] and not host[end + 1 :].startswith(":")):
return False # unclosed bracket or junk after ] (e.g. [::1]evil)
host = host[1:end]
elif host.count(":") == 1: # host:port
host = host.split(":", 1)[0]
host = host.lower().rstrip(".")
return host == "localhost" or _is_loopback_ip(host)
def _is_local_bootstrap_request(request: Request) -> bool:
"""Allow bootstrap injection only through a direct loopback authority."""
client = request.client
if client is None or not _is_loopback_ip(client.host):
return False
if any(request.headers.get(h) is not None for h in _PROXIED_CLIENT_HEADERS):
return False
return _host_header_is_loopback(request.headers.get("host"))
def _is_same_origin_request(request: Request) -> bool:
"""True when Origin is missing or matches request's scheme://host:port.
@ -1390,6 +1460,17 @@ def _is_same_origin_request(request: Request) -> bool:
return origin_canon == self_canon
def _should_inject_bootstrap(request: Request) -> bool:
"""Whether to embed the seeded bootstrap password in index.html."""
if not _is_same_origin_request(request):
return False
if _IS_COLAB:
# Single-user notebook proxy: allow autofill, but never a public
# shareable tunnel (a Colab Cloudflare link sets cf-connecting-ip).
return request.headers.get("cf-connecting-ip") is None
return _is_local_bootstrap_request(request)
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@ -1402,8 +1483,10 @@ def setup_frontend(app: FastAPI, build_path: Path):
def _build_index_response(request: Request) -> Response:
content = (build_path / "index.html").read_bytes()
content = _strip_crossorigin(content)
# Bootstrap pw is same-origin only; Vary: Origin keeps caches honest.
if _is_same_origin_request(request):
# Bootstrap pw goes only to a same-origin, direct-loopback client (or
# Colab's single-user notebook proxy): a wildcard bind must not serve it
# in-page to a LAN or proxied peer. Vary: Origin keeps caches honest.
if _should_inject_bootstrap(request):
content, nonce = _inject_bootstrap(content, app)
else:
nonce = None

10
studio/backend/models/auth.py
View file

@ -7,6 +7,8 @@ from typing import Optional
from pydantic import BaseModel, Field
from auth.storage import MIN_PASSWORD_LENGTH
class AuthLoginRequest(BaseModel):
"""Login payload: username/password to obtain a JWT."""
@ -45,10 +47,14 @@ class ChangePasswordRequest(BaseModel):
"""Change the current user's password, typically on first login."""
current_password: str = Field(
..., min_length = 8, description = "Existing password for the authenticated user"
...,
min_length = MIN_PASSWORD_LENGTH,
description = "Existing password for the authenticated user",
)
new_password: str = Field(
..., min_length = 8, description = "Replacement password (minimum 8 characters)"
...,
min_length = MIN_PASSWORD_LENGTH,
description = f"Replacement password (minimum {MIN_PASSWORD_LENGTH} characters)",
)

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

@ -140,6 +140,27 @@ class ValidateModelRequest(BaseModel):
)
class TransformersUpgradeInfo(BaseModel):
"""A model architecture no installed transformers ships, but a newer release does."""
model_type: str = Field(
..., description = "config.json model_type unknown to every installed transformers"
)
pypi_version: Optional[str] = Field(
None, description = "Latest transformers release on PyPI at check time"
)
supported_in_pypi: bool = Field(
False,
description = "True if the latest PyPI release ships this model_type; Studio can "
"install it into a persistent sidecar after user consent.",
)
supported_in_main: bool = Field(
False,
description = "True if transformers GitHub main ships this model_type (dev-only; "
"not installable through Studio yet).",
)
class ValidateModelResponse(BaseModel):
"""Result of model validation.
@ -167,6 +188,48 @@ class ValidateModelResponse(BaseModel):
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,
description = "True when the model's architecture is unknown to every installed "
"transformers but a newer transformers ships it; the UI should offer the "
"install-latest-transformers consent dialog (or the dev-only notice).",
)
transformers_upgrade: Optional[TransformersUpgradeInfo] = Field(
None,
description = "Details for the transformers-upgrade dialog; set only when "
"requires_transformers_upgrade is true.",
)
class InstallLatestTransformersRequest(BaseModel):
"""Consented request to install the latest transformers release into a sidecar."""
version: str = Field(
...,
min_length = 1,
max_length = 64,
description = "Exact transformers version to install; must match the current "
"latest PyPI release reported by /validate.",
)
class InstallLatestTransformersResponse(BaseModel):
"""Result of the consented latest-transformers sidecar install."""
success: bool = Field(..., description = "Whether the sidecar was provisioned")
version: str = Field(..., description = "The requested transformers version")
message: str = Field(..., description = "Human-readable result")
model_unloaded: bool = Field(
False,
description = "Whether the active chat model was unloaded before the swap "
"(reported even on failure, so the client can restore its state)",
)
latest_version: Optional[str] = Field(
None,
description = "On a version-mismatch failure: the release that superseded "
"the requested one, so the client can retry with it",
)
class GenerateRequest(BaseModel):
@ -632,6 +695,23 @@ class ThinkingConfig(BaseModel):
type: Literal["disabled", "enabled"] = "disabled"
# Recognized permission_mode values. The field accepts a plain string rather than
# a Literal so an unrecognized value from a newer UI/client degrades to the
# safest gate ("ask") instead of a 422; the tool loops apply the same unknown ->
# ask fallback, so normalizing here keeps that forward-compat path reachable at
# the API boundary. None stays unset ("behaves as 'ask'" without self-enabling
# the confirm gate).
_KNOWN_PERMISSION_MODES = ("ask", "auto", "off", "full")
def _normalize_permission_mode(value: Any) -> Any:
if value is None:
return None
if value not in _KNOWN_PERMISSION_MODES:
return "ask"
return value
class ChatCompletionRequest(BaseModel):
"""OpenAI-compatible chat completion request.
@ -777,6 +857,19 @@ class ChatCompletionRequest(BaseModel):
False,
description = "[x-unsloth] Bypass Permissions: when true, skip the tool-call confirmation gate AND disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits). Secret env vars are still stripped. Takes precedence over confirm_tool_calls.",
)
permission_mode: Optional[str] = Field(
None,
description = (
"[x-unsloth] Permission level for local tool calls. 'ask' pauses every "
"call for approval; 'ask'/'auto' enable the confirmation gate on their "
"own (needs a streaming request to deliver prompts). 'auto' ('Approve for "
"me') only pauses calls detected as potentially unsafe (state-mutating "
"terminal/python/MCP calls); read-only calls run immediately, and the "
"sandbox stays on. 'full' is equivalent to bypass_permissions=true (no "
"confirmation, no sandbox). Unset behaves as 'ask'. An unrecognized value "
"(e.g. from a newer client) is treated as 'ask'."
),
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@ -815,6 +908,10 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
)
thread_id: Optional[str] = Field(
None,
description = "[x-unsloth] Conversation ID for scoping stateful tool sessions (e.g. stdio MCP); stays per-thread where session_id may be shared project-wide.",
)
rag_scope: Optional[dict] = Field(
None,
description = (
@ -1036,6 +1133,52 @@ class ChatCompletionRequest(BaseModel):
self.enable_thinking = self.thinking.type == "enabled"
return self
@field_validator("permission_mode", mode = "before")
@classmethod
def _coerce_permission_mode(cls, value: Any) -> Any:
# Accept any string so an unknown mode degrades to 'ask' instead of a
# 422; mirrors the tool loops' unknown -> ask fallback.
return _normalize_permission_mode(value)
@model_validator(mode = "after")
def _fold_full_permission_into_bypass(self) -> "ChatCompletionRequest":
"""permission_mode='full' is the documented equivalent of
bypass_permissions=true, so fold it in before any route guard reads
the flag (else a full request would trip the confirm-gate rejections)."""
if self.permission_mode == "full":
self.bypass_permissions = True
elif self.bypass_permissions:
# Legacy bypass callers map onto Full access (mirrors the tool loop).
self.permission_mode = "full"
elif self.permission_mode == "off":
# "Off" never prompts, so route guards must see confirm disabled.
self.confirm_tool_calls = False
elif (
self.permission_mode == "ask"
and self.confirm_tool_calls is None
and not (self.provider_id or self.provider_type)
and (self.enable_tools is True or bool(self.mcp_enabled))
):
# "Ask" gates every call, so a direct API caller that omits the legacy
# confirm flag must still hit the confirmation gate for Studio's own
# tool loop. An explicit confirm_tool_calls=False wins over the mode
# (mirrors _permission_mode_confirm and the Anthropic pre-switch guard),
# so only self-enable when the flag is unset. Only self-enable when that
# loop is actually requested
# (enable_tools / mcp_enabled) -- the router enters the loop on those
# signals, not on enabled_tools alone (which merely filters which tools
# run). A plain client-tool passthrough (client-supplied `tools` that
# Studio does not execute) must route verbatim, and external-provider
# routing rejects confirm_tool_calls with tools, so skip the fold there.
#
# "auto" is deliberately NOT folded: it only prompts for a call the
# classifier flags, so leaving confirm_tool_calls unset lets the route's
# _confirm_gate_needs_stream apply the safe-only exception (a safe-only
# auto selection needs no stream) instead of an explicit-confirm forcing
# stream=true. The mode still drives the loop's per-call gate.
self.confirm_tool_calls = True
return self
class ToolConfirmRequest(BaseModel):
session_id: Optional[str] = None
@ -1533,12 +1676,41 @@ class AnthropicToolResultBlock(BaseModel):
tool_use_id: str
content: Union[str, list] = ""
@field_validator("content", mode = "before")
@classmethod
def _coerce_null_content(cls, v):
# Some clients send null content for an empty tool result; the str|list
# union would 400 on it, so treat null as "".
return "" if v is None else v
# Block types the converter translates explicitly. Anything else (thinking /
# redacted_thinking, a provider block a resumed session replays, or a future type)
# is accepted as an unknown block and dropped by the converter, rather than 400-ing
# the whole request on strict validation.
_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"})
class AnthropicUnknownBlock(BaseModel):
type: str
model_config = {"extra": "allow"}
@field_validator("type")
@classmethod
def _only_unknown_types(cls, v):
# Known types parse as their typed models above (so a malformed known block
# still fails cleanly); this fallback only catches the rest.
if v in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError("known block type handled by its typed model")
return v
AnthropicContentBlock = Union[
AnthropicTextBlock,
AnthropicImageBlock,
AnthropicToolUseBlock,
AnthropicToolResultBlock,
AnthropicUnknownBlock,
]
@ -1583,6 +1755,40 @@ class AnthropicMessage(BaseModel):
role: Literal["user", "assistant"]
content: Union[str, list[AnthropicContentBlock]]
@model_validator(mode = "before")
@classmethod
def _normalize_content(cls, data):
# Role-aware leniency that never silently drops real user input:
# - assistant: a resumed tool-only turn's null content -> "" (str|list would
# 400 on null; "" keeps the converter's `for block in content` safe).
# Unknown blocks (thinking / future types) validate via
# AnthropicUnknownBlock and are dropped by the converter.
# - user: keep strict. Null user content stays None so str|list rejects it
# (400) rather than forwarding an empty prompt; and reject block types the
# converter cannot translate, since it silently skips unknown user blocks
# -- a user turn made only of them would validate yet send no content
# (silent data loss).
if not isinstance(data, dict):
return data
content = data.get("content")
if data.get("role") == "assistant":
# Coerce only an explicit null (resumed tool-only turn). A missing
# content key stays malformed so the required-field check still 400s.
if "content" in data and content is None:
return {**data, "content": ""}
return data
if isinstance(content, list):
for block in content:
btype = (
block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
)
# Guard the value: a non-string type is unsupported too, and a
# membership test on an unhashable value would raise TypeError
# (escaping as a 500 instead of a clean 400).
if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError(f"unsupported content block type {btype!r} in a user message")
return data
class AnthropicTool(BaseModel):
# Client tools have input_schema; server tools may only have type/name.
@ -1619,11 +1825,19 @@ class AnthropicMessagesRequest(BaseModel):
enable_tools: Optional[bool] = None
enabled_tools: Optional[list[str]] = None
session_id: Optional[str] = None
thread_id: Optional[str] = Field(
None,
description = "[x-unsloth] Conversation ID for scoping stateful tool sessions (e.g. stdio MCP); stays per-thread where session_id may be shared project-wide.",
)
cancel_id: Optional[str] = None
bypass_permissions: Optional[bool] = Field(
False,
description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
)
permission_mode: Optional[str] = Field(
None,
description = "[x-unsloth] Permission level for local tool calls: 'ask' pauses every call, 'auto' only pauses calls detected as potentially unsafe, 'off' never pauses (sandbox stays on), 'full' equals bypass_permissions=true. Unset behaves as 'ask'; an unrecognized value (e.g. from a newer client) is treated as 'ask'. Declared explicitly so omitted requests default to None instead of raising AttributeError.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).",
@ -1665,6 +1879,27 @@ class AnthropicMessagesRequest(BaseModel):
normalized["system"] = _merge_anthropic_system(normalized.get("system"), system_additions)
return normalized
@field_validator("permission_mode", mode = "before")
@classmethod
def _coerce_permission_mode(cls, value: Any) -> Any:
# Accept any string so an unknown mode degrades to 'ask' instead of a
# 422; mirrors the tool loops' unknown -> ask fallback.
return _normalize_permission_mode(value)
@model_validator(mode = "after")
def _fold_full_permission_into_bypass(self) -> "AnthropicMessagesRequest":
"""permission_mode='full' equals bypass_permissions=true (mirrors the
Chat Completions request)."""
if self.permission_mode == "full":
self.bypass_permissions = True
elif self.bypass_permissions:
# Legacy bypass callers map onto Full access (mirrors the tool loop).
self.permission_mode = "full"
elif self.permission_mode == "off":
# "Off" never prompts, so route guards must see confirm disabled.
self.confirm_tool_calls = False
return self
# ── Response models ────────────────────────────────────────────

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

@ -500,8 +500,9 @@ async def change_password(
detail = "New password must be different from the current password",
)
storage.update_password(current_subject, payload.new_password)
storage.revoke_user_refresh_tokens(current_subject)
# Single transaction: a separate refresh-token purge could fail after the
# password commit, leaving pre-change tokens able to mint access tokens.
storage.update_password(current_subject, payload.new_password, revoke_refresh_tokens = True)
try:
request.app.state.bootstrap_password = None
except AttributeError:

37
studio/backend/routes/data_recipe/seed.py
View file

@ -10,6 +10,7 @@ import binascii
import json
import os
import re
import shutil
from itertools import islice
from pathlib import Path
from typing import Any
@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"}
SEED_UPLOAD_DIR = seed_uploads_root()
UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root()
_SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$")
# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...)
# never match: those directories can be shared by several recipes.
_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$")
def _validate_safe_id(value: str, label: str) -> str:
@ -580,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str):
return {"status": "ok"}
@router.delete("/seed/unstructured-block/{block_id}")
async def remove_unstructured_block(block_id: str):
"""Delete a block's upload directory; files on disk still count toward its quota.
Only uid-namespaced directories may be bulk-deleted: they have exactly one
owning block. Legacy node-id directories (n1, ...) can be shared by other
recipes, so they are managed file-by-file instead.
"""
_validate_safe_id(block_id, "block_id")
if not _UPLOAD_UID_RE.match(block_id):
raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted")
block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve()
if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()):
raise HTTPException(400, "Invalid block_id: outside upload root")
if not block_dir.exists():
return {"status": "ok", "deleted": False}
try:
shutil.rmtree(block_dir)
except OSError as exc:
raise log_and_http_error(
exc,
500,
"failed to delete uploaded files",
event = "data_recipe.seed.unstructured_block_delete_failed",
log = logger,
) from exc
if block_dir.exists():
raise HTTPException(500, "failed to delete uploaded files")
return {"status": "ok", "deleted": True}
@router.post("/seed/inspect-upload", response_model = SeedInspectResponse)
def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse:
if payload.file_ids is not None:

35
studio/backend/routes/export.py
View file

@ -51,7 +51,17 @@ def _ensure_export_supported() -> None:
Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints
(scan/status/logs) are intentionally NOT gated so the Export page can still render the reason.
Also refuses (409) while a latest-transformers install is swapping .venv_t5_latest: an
export worker spawned mid-swap could activate a half-replaced sidecar.
"""
from utils.transformers_latest import is_install_in_progress
if is_install_in_progress():
raise HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
from utils.hardware import export_capability
cap = export_capability()
@ -97,6 +107,11 @@ async def load_checkpoint(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error loading checkpoint: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -308,6 +323,11 @@ async def export_merged_model(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting merged model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -347,6 +367,11 @@ async def export_base_model(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting base model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -388,6 +413,11 @@ async def export_gguf(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting GGUF model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -428,6 +458,11 @@ async def export_lora_adapter(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True)
raise HTTPException(
status_code = 500,

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

File diff suppressed because it is too large Load diff

28
studio/backend/routes/llama.py
View file

@ -14,14 +14,17 @@ never blocks on a missing marker / offline GitHub.
from __future__ import annotations
import asyncio
import threading
from typing import Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject
from loggers import get_logger
from utils.llama_cpp_update import get_update_status, start_update
logger = get_logger(__name__)
router = APIRouter()
@ -69,6 +72,27 @@ class LlamaUpdateActionResponse(BaseModel):
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
_llama_update_lock = threading.Lock()
_last_llama_update_step = -1
def _log_llama_update_progress(job: LlamaUpdateJob) -> None:
"""One llama_update_progress line per 10% step so a prebuilt update reports
progress without a line per poll. Resyncs when a new update starts."""
global _last_llama_update_step
if job.state != "running" or job.progress is None:
return
step = int(max(0.0, min(float(job.progress), 1.0)) * 10)
with _llama_update_lock:
prev = _last_llama_update_step
if step == prev:
return
_last_llama_update_step = step
if step < prev:
return # new update; resync without logging
logger.info("llama_update_progress", to_tag = job.to_tag or "", percent = step * 10)
@router.get("/update-status", response_model = LlamaUpdateStatusResponse)
async def llama_update_status(
force_refresh: bool = Query(
@ -78,7 +102,9 @@ async def llama_update_status(
) -> LlamaUpdateStatusResponse:
# Off the event loop: detection may probe the host and read GitHub.
status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh)
return LlamaUpdateStatusResponse(**status)
resp = LlamaUpdateStatusResponse(**status)
_log_llama_update_progress(resp.job)
return resp
@router.post("/update", response_model = LlamaUpdateActionResponse)

15
studio/backend/routes/mcp_servers.py
View file

@ -1,6 +1,7 @@
# 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 asyncio
import json
import uuid
from urllib.parse import urlparse
@ -13,6 +14,7 @@ from core.inference.mcp_client import (
TOOL_CACHE_INVALIDATING_FIELDS,
cache_tools,
clear_oauth_tokens_async,
close_stdio_sessions,
invalidate_tool_cache,
is_stdio,
list_tools_async,
@ -206,11 +208,15 @@ async def update_mcp_server(
):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.update_server(server_id, changes)
# A new endpoint/auth makes cached tools wrong and disabling makes them
# unreachable, so drop them and let the next send re-probe; a rename
# leaves them valid.
if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS:
# A new endpoint/auth makes cached tools wrong and disabling makes them unreachable, so drop
# them and let the next send re-probe; a rename leaves them valid. Live stdio sessions for the
# old endpoint close too. Gate on a real value change, not mere presence: the edit dialog
# resends url/headers/oauth unchanged on a rename, which must not drop the session.
if any(changes[k] != old.get(k) for k in changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS):
invalidate_tool_cache(server_id)
# Narrow to this row's env: another server row sharing the command but
# with a different env keeps its live sessions.
await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old))
return _row_to_response(mcp_servers_db.get_server(server_id))
@ -223,6 +229,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.delete_server(server_id)
invalidate_tool_cache(server_id)
await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old))
@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)

113
studio/backend/routes/models.py
View file

@ -1188,7 +1188,9 @@ def _looks_like_model_dir(directory: Path) -> bool:
return False
def _build_browse_allowlist() -> list[Path]:
def _build_browse_allowlist(
media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None
) -> list[Path]:
"""Return the root directories the folder browser may walk.
The same list seeds the sidebar suggestion chips, so chip targets are
@ -1196,13 +1198,20 @@ def _build_browse_allowlist() -> list[Path]:
outputs/exports/studio root, registered scan folders, and well-known
local-LLM dirs (LM Studio, Ollama, ``~/models``); each added only if
it resolves to a real directory.
*media_roots* / *drive_roots* let the caller pass already-probed
removable-media and Windows drive roots so they aren't scanned again (a
disconnected mapped drive can make each probe slow); probed here when ``None``.
"""
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from storage.studio_db import list_scan_folders
candidates: list[Path] = []
@ -1218,7 +1227,13 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
if media_roots is None:
media_roots = linux_run_media_mount_roots()
if drive_roots is None:
drive_roots = windows_drive_roots()
for p in media_roots:
_add(p)
for p in drive_roots:
_add(p)
_add(_resolve_hf_cache_dir())
try:
@ -1269,19 +1284,43 @@ def _build_browse_allowlist() -> list[Path]:
def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
"""True if *target* equals or descends from any allowed root.
Uses ``os.path.realpath`` so symlinks can't escape the sandbox.
Uses ``os.path.realpath`` (symlinks can't escape the sandbox) and
``os.path.commonpath`` for a component-wise containment test, so a string
prefix like ``/home/u`` never matches a sibling ``/home/user2`` while a
drive root ``D:\\`` still contains ``D:\\models``. A Windows drive root
authorizes its descendants, but a bare POSIX root ``/`` must NOT, else one
``/`` allowlist entry would authorize every absolute path. ``normcase`` keeps
the drive-letter comparison case-insensitive, matching the hub browser.
"""
try:
target_real = os.path.realpath(str(target))
target_real = os.path.normcase(os.path.realpath(str(target)))
except OSError:
return False
for root in allowed_roots:
try:
root_real = os.path.realpath(str(root))
root_real = os.path.normcase(os.path.realpath(str(root)))
except OSError:
continue
if target_real == root_real or target_real.startswith(root_real + os.sep):
if target_real == root_real:
return True
drive, tail = os.path.splitdrive(root_real)
if os.path.dirname(root_real) == root_real and not drive:
# Bare POSIX filesystem root ("/"): equality above is the only
# match; do not let it authorize arbitrary descendants.
continue
if drive.startswith(("\\\\", "//")) and not tail:
# Bare UNC share root (\\server\share): os.path.commonpath raises
# "can't mix absolute and relative" on it, so authorize its
# descendants with a boundary-safe prefix test (normcase applied).
if target_real.startswith(root_real.rstrip("\\/") + os.sep):
return True
continue
try:
if os.path.commonpath([target_real, root_real]) == root_real:
return True
except ValueError:
# Different drives / mixed absolute-relative: not contained.
continue
return False
@ -1339,7 +1378,10 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]:
def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path:
"""Resolve a requested browse path by walking from trusted allowlist roots."""
from storage.studio_db import contains_sensitive_path_component
from storage.studio_db import (
contains_sensitive_path_component,
is_denied_system_path,
)
requested_path = _normalize_browse_request_path(path)
resolved_roots: list[Path] = []
@ -1396,6 +1438,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if is_denied_system_path(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
@ -1403,6 +1450,13 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
# Zero-component case: the requested path IS an allowlist root
# (e.g. a legacy-registered "/" or a Windows drive root).
if is_denied_system_path(str(current)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -1420,8 +1474,12 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
)
# Sync (def, not async) so FastAPI runs the blocking filesystem I/O (drive
# probes, iterdir, realpath) in the threadpool: a disconnected mapped drive can
# make the probe wait out its timeout, which on the event loop would stall every
# other request. Matches the hub browse endpoint.
@router.get("/browse-folders", response_model = BrowseFoldersResponse)
async def browse_folders(
def browse_folders(
path: Optional[str] = Query(
None,
description = (
@ -1450,11 +1508,22 @@ async def browse_folders(
then hidden (if ``show_hidden=true``).
"""
from utils.paths import hf_default_cache_dir, well_known_model_dirs
from utils.paths.external_media import linux_run_media_mount_roots
from storage.studio_db import contains_sensitive_path_component, list_scan_folders
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from storage.studio_db import (
contains_sensitive_path_component,
is_denied_system_path,
list_scan_folders,
)
# Probe removable-media and Windows drive roots once; the allowlist and
# chips reuse the result so a disconnected mapped drive isn't scanned twice.
media_roots = linux_run_media_mount_roots()
drive_roots = windows_drive_roots()
# Build once; the sandbox check and suggestion chips share it.
allowed_roots = _build_browse_allowlist()
allowed_roots = _build_browse_allowlist(media_roots, drive_roots)
try:
target = _resolve_browse_target(path, allowed_roots)
@ -1506,6 +1575,15 @@ async def browse_folders(
continue
if contains_sensitive_path_component(name):
continue
# Hide denied system dirs (C:\Windows, /etc, ...) so they don't
# render as clickable rows that then 403 on descent. Resolve first
# so a symlink/junction into a denied dir is hidden too, not just a literal name.
try:
resolved_child = os.path.realpath(str(child))
except (OSError, ValueError):
resolved_child = str(child)
if is_denied_system_path(resolved_child):
continue
entries.append(
BrowseEntry(
name = name,
@ -1553,13 +1631,22 @@ async def browse_folders(
return
if resolved in seen_sug:
return
# Drop a denied system dir (e.g. a stale scan-folder row) so it never
# becomes a chip that 403s on click. Drive roots stay: only their
# system subdirectories are denied, not the root itself.
if is_denied_system_path(resolved):
return
if _safe_is_dir(resolved):
seen_sug.add(resolved)
suggestions.append(resolved)
# Home first -- the safe fallback when everything else is cold.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
# Reuse the roots probed for the allowlist above (no second drive scan).
for p in media_roots:
_add_sug(p)
# Windows drive roots so the user can hop between C:, D:, E: ...
for p in drive_roots:
_add_sug(p)
# The HF cache root the process is actually using.
try:

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

@ -1,6 +1,7 @@
# 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 re
from typing import Literal, Optional
from urllib.parse import unquote, urlsplit
@ -32,6 +33,7 @@ from utils.helper_precache_settings import (
helper_model_disabled_by_env,
set_helper_precache_enabled,
)
from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents
from utils.openai_auto_switch_settings import (
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
@ -174,6 +176,19 @@ def update_helper_precache(
return _helper_precache_response(enabled)
class CodingAgentsResponse(BaseModel):
# All agents `unsloth start` supports, in the CLI's declared order.
agents: tuple[str, ...] = CODING_AGENTS
# Subset of `agents` whose CLI binary was found on PATH; the frontend uses
# this to default the API-keys panel to a command the user can run as-is.
detected: list[str]
@router.get("/coding-agents", response_model = CodingAgentsResponse)
def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse:
return CodingAgentsResponse(detected = detect_installed_coding_agents())
@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse)
def get_openai_auto_switch(
current_subject: str = Depends(get_current_subject),
@ -537,6 +552,7 @@ class PersonalizationProfile(BaseModel):
nickname: str = Field("", max_length = 200)
avatarDataUrl: Optional[str] = Field(None, max_length = MAX_AVATAR_DATA_URL_BYTES)
avatarShape: Literal["circle", "rounded"] = "circle"
showGreetingSloth: bool = True
@field_validator("avatarDataUrl")
@classmethod
@ -548,11 +564,179 @@ class PersonalizationProfile(BaseModel):
return value
class PersonalizationCustomColors(BaseModel):
model_config = ConfigDict(extra = "ignore")
accent: Optional[str] = Field(None, pattern = r"^#[0-9a-fA-F]{6}$")
background: Optional[str] = Field(None, pattern = r"^#[0-9a-fA-F]{6}$")
foreground: Optional[str] = Field(None, pattern = r"^#[0-9a-fA-F]{6}$")
class PersonalizationCustomColorModes(BaseModel):
model_config = ConfigDict(extra = "ignore")
light: PersonalizationCustomColors = Field(default_factory = PersonalizationCustomColors)
dark: PersonalizationCustomColors = Field(default_factory = PersonalizationCustomColors)
MAX_IMPORTED_FONTS = 3
# ~1.5 MB font file as base64; matches MAX_IMPORTED_FONT_DATA_URL_LENGTH in
# the frontend appearance-custom-store.
MAX_FONT_DATA_URL_LENGTH = 2_200_000
# Aggregate cap across all imported fonts; matches
# MAX_TOTAL_IMPORTED_FONT_DATA_URL_LENGTH in the frontend so a synced payload
# always fits the browser's localStorage quota.
MAX_TOTAL_FONT_DATA_URL_LENGTH = 4_400_000
# Characters that could terminate a CSS declaration, escape the quoted
# font-family value (backslash), or smuggle extra fallbacks/comments (comma,
# slash) if a stored name ever reached a stylesheet. The server is the
# authoritative gate; the frontend strips the same set before use.
_FONT_NAME_FORBIDDEN = set(";{}()<>\"'\\/,`")
def _check_font_name(value: str) -> str:
if any(c in _FONT_NAME_FORBIDDEN or ord(c) < 0x20 for c in value):
raise ValueError("Font name contains invalid characters.")
return value
# Matches FONT_DATA_URL_PATTERN in the frontend appearance-custom-store.
_FONT_DATA_URL_PATTERN = re.compile(
r"^data:(?:font/(?:woff2?|ttf|otf|sfnt)"
r"|application/(?:octet-stream|x-font-\w+|font-\w+));base64,[A-Za-z0-9+/=]+$"
)
class PersonalizationImportedFont(BaseModel):
model_config = ConfigDict(extra = "ignore")
name: str = Field(..., min_length = 1, max_length = 100)
dataUrl: str = Field(..., max_length = MAX_FONT_DATA_URL_LENGTH)
@field_validator("name")
@classmethod
def _validate_font_name(cls, value: str) -> str:
return _check_font_name(value)
@field_validator("dataUrl")
@classmethod
def _validate_font_data_url(cls, value: str) -> str:
# fullmatch, not match: re's ``$`` also matches just before a trailing
# newline, so ``match`` would accept "data:font/woff2;base64,AAAA\n",
# which the frontend's JS pattern (``$`` = end of string) rejects.
if not _FONT_DATA_URL_PATTERN.fullmatch(value):
raise ValueError("dataUrl must be a base64 font data URL.")
return value
# Optional user-menu items; the boolean is each id's default visibility.
# Settings-tab shortcuts ship hidden.
SIDEBAR_MENU_ITEM_DEFAULTS = {
"api": True,
"darkMode": True,
"guidedTour": True,
"profile": False,
"appearance": False,
"resources": False,
"chat": False,
"connections": False,
}
# The sidebarMenu validator below dedupes ids and re-fills any missing ones, so
# the stored list is always exactly one entry per id. Cap the *incoming* list at
# a generous multiple rather than len(defaults): a stale or duplicated payload
# (more items than distinct ids) must reach the validator so it can normalize,
# instead of being rejected by the length constraint before dedupe runs. A
# pathologically long list is still refused.
MAX_SIDEBAR_MENU_INPUT_ITEMS = 4 * len(SIDEBAR_MENU_ITEM_DEFAULTS)
class PersonalizationSidebarMenuItem(BaseModel):
model_config = ConfigDict(extra = "ignore")
id: Literal[
"api",
"darkMode",
"guidedTour",
"profile",
"appearance",
"resources",
"chat",
"connections",
]
visible: bool = True
def _default_sidebar_menu() -> "list[PersonalizationSidebarMenuItem]":
return [
PersonalizationSidebarMenuItem(id = item_id, visible = visible)
for item_id, visible in SIDEBAR_MENU_ITEM_DEFAULTS.items()
]
class PersonalizationCustomization(BaseModel):
model_config = ConfigDict(extra = "ignore")
colors: PersonalizationCustomColorModes = Field(default_factory = PersonalizationCustomColorModes)
uiFont: Optional[str] = Field(None, max_length = 200)
headingFont: Optional[str] = Field(None, max_length = 200)
chatFont: Optional[str] = Field(None, max_length = 200)
codeFont: Optional[str] = Field(None, max_length = 200)
importedFonts: list[PersonalizationImportedFont] = Field(
default_factory = list, max_length = MAX_IMPORTED_FONTS
)
@field_validator("importedFonts")
@classmethod
def _validate_total_font_size(
cls, value: list[PersonalizationImportedFont]
) -> list[PersonalizationImportedFont]:
if sum(len(f.dataUrl) for f in value) > MAX_TOTAL_FONT_DATA_URL_LENGTH:
raise ValueError("Imported fonts exceed the total size limit.")
return value
@field_validator("uiFont", "headingFont", "chatFont", "codeFont")
@classmethod
def _validate_selected_fonts(cls, value: Optional[str]) -> Optional[str]:
# Selected font names reach CSS the same way imported names do.
return value if value is None else _check_font_name(value)
uiFontSize: Optional[int] = Field(None, ge = 12, le = 20)
codeFontSize: Optional[int] = Field(None, ge = 10, le = 20)
contrast: int = Field(50, ge = 0, le = 100)
pointerCursors: bool = False
reduceMotion: Literal["system", "on", "off"] = "system"
fontSmoothing: bool = True
sidebarMenu: list[PersonalizationSidebarMenuItem] = Field(
default_factory = _default_sidebar_menu,
max_length = MAX_SIDEBAR_MENU_INPUT_ITEMS,
)
@field_validator("sidebarMenu")
@classmethod
def _validate_sidebar_menu(
cls, value: list[PersonalizationSidebarMenuItem]
) -> list[PersonalizationSidebarMenuItem]:
# Drop duplicate ids (keep the first) and re-append any missing ids so
# the stored list always covers every optional menu item exactly once.
seen: set[str] = set()
items = [item for item in value if not (item.id in seen or seen.add(item.id))]
for item_id, visible in SIDEBAR_MENU_ITEM_DEFAULTS.items():
if item_id not in seen:
items.append(PersonalizationSidebarMenuItem(id = item_id, visible = visible))
return items
class PersonalizationAppearance(BaseModel):
model_config = ConfigDict(extra = "ignore")
theme: Literal["light", "dark", "system"] = "system"
palette: Literal["standard", "classic", "minimal"] = "standard"
language: Optional[str] = Field(None, max_length = 20)
customization: PersonalizationCustomization = Field(
default_factory = PersonalizationCustomization
)
class PersonalizationPayload(BaseModel):
@ -565,6 +749,11 @@ class PersonalizationPayload(BaseModel):
class PersonalizationResponse(PersonalizationPayload):
saved: bool = False
# False when the stored record predates a field, so the client keeps local
# overrides instead of treating a server-filled default as an explicit value.
customizationSaved: bool = False
paletteSaved: bool = False
greetingSlothSaved: bool = False
@router.get("/personalization", response_model = PersonalizationResponse)
@ -574,15 +763,38 @@ def get_personalization_settings(
stored = get_personalization()
response = PersonalizationResponse.model_validate(stored or {})
response.saved = bool(stored)
appearance = stored.get("appearance") if isinstance(stored, dict) else None
profile = stored.get("profile") if isinstance(stored, dict) else None
response.customizationSaved = isinstance(appearance, dict) and "customization" in appearance
response.paletteSaved = isinstance(appearance, dict) and "palette" in appearance
response.greetingSlothSaved = isinstance(profile, dict) and "showGreetingSloth" in profile
return response
def _merge_personalization(base: dict, overlay: dict) -> dict:
# Recursively overlay only the request's set fields onto the stored record,
# so a stale client that omits newer keys (palette, customization) does not
# materialize their defaults and defeat the *Saved legacy detection.
merged = dict(base)
for key, value in overlay.items():
existing = merged.get(key)
if isinstance(value, dict) and isinstance(existing, dict):
merged[key] = _merge_personalization(existing, value)
else:
merged[key] = value
return merged
@router.put("/personalization", response_model = PersonalizationPayload)
def update_personalization_settings(
payload: PersonalizationPayload, current_subject: str = Depends(get_current_subject)
) -> PersonalizationPayload:
try:
set_personalization(payload.model_dump())
# exclude_unset so absent fields are not persisted as defaults; merge so
# fields the request omits keep whatever the record already stored.
incoming = payload.model_dump(exclude_unset = True)
merged = _merge_personalization(get_personalization(), incoming)
set_personalization(merged)
except ValueError as exc:
raise log_and_http_error(
exc,
@ -591,4 +803,6 @@ def update_personalization_settings(
event = "settings.update_personalization_failed",
log = logger,
) from exc
return payload
# Return the stored record, not the defaults-filled request, so the response
# matches storage (and the next GET) for fields the client omitted.
return PersonalizationPayload.model_validate(merged)

45
studio/backend/routes/training.py
View file

@ -146,6 +146,16 @@ async def start_training(
# No in-process ensure_transformers_version(): the subprocess
# (worker.py) activates the correct version before importing ML libs.
# A consented latest-transformers install stage-and-swaps .venv_t5_latest;
# a worker spawned mid-swap could activate a half-replaced sidecar.
from utils.transformers_latest import is_install_in_progress
if is_install_in_progress():
raise HTTPException(
status_code = 409,
detail = ("A transformers installation is in progress. Retry when it completes."),
)
backend = get_training_backend()
# S3 dataset loading needs the optional boto3 dependency. Reject early
@ -341,6 +351,24 @@ async def start_training(
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
}
# Latest-sidecar models size and train 16-bit (same flip as chat load):
# 4-bit is disabled for brand-new architectures, so VRAM coexistence
# checks must not underestimate against a load the worker will refuse.
if training_kwargs["load_in_4bit"]:
from utils.transformers_version import latest_tier_active_for
if await asyncio.to_thread(
latest_tier_active_for,
training_kwargs["model_name"],
training_kwargs["hf_token"] or None,
):
training_kwargs["load_in_4bit"] = False
logger.info(
"Latest-transformers sidecar active for %s - sizing and "
"training in 16-bit (4-bit is disabled for brand-new "
"architectures)",
training_kwargs["model_name"],
)
# Training page has no trust_remote_code toggle, so honor the YAML default
# -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
# local path or a name merely starting with "unsloth/".
@ -426,9 +454,16 @@ async def start_training(
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
from utils.transformers_version import SidecarSwapInProgress
try:
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
except SidecarSwapInProgress as exc:
# Expected loss of the race against a sidecar install: a retryable
# 409 matching the route-entry guard, not an internal error.
raise HTTPException(status_code = 409, detail = str(exc))
if not success:
progress_error = backend.trainer.training_progress.error
@ -698,7 +733,9 @@ async def stream_training_progress(
if last_event_id is not None:
try:
resume_from_step = int(last_event_id)
logger.info(f"SSE reconnect: resuming from step {resume_from_step}")
# Fires on every reconnect (each tab switch); the meaningful signal is
# the "replayed N missed steps" line below, logged only when N > 0.
logger.debug(f"SSE reconnect: resuming from step {resume_from_step}")
except ValueError:
logger.warning(f"Invalid Last-Event-ID: {last_event_id}")

317
studio/backend/run.py
View file

@ -10,7 +10,7 @@ import os
import sys
import time
from pathlib import Path
from typing import Optional
from typing import Optional, Tuple
def _fix_torch_cuda_ld_path():
@ -616,24 +616,28 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
elif not _cloudflare_flag:
elif _cloudflare_flag is False or _cloudflare_flag is None:
# None = off by default (no flag); False = explicit --no-cloudflare.
_reason = "default" if _cloudflare_flag is None else "--no-cloudflare"
if _public_reachable is True:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). The raw port is still "
f" Cloudflare tunnel: OFF ({_reason}). The raw port is still "
"reachable from the public internet (see the reachability check above): "
"--no-cloudflare disables only the Cloudflare link, not the public bind.",
"pass --cloudflare to also expose a public Cloudflare HTTPS link, or "
f"bind {loopback_host} to keep Studio private.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). Studio is reachable on your "
"local network only. Omit --no-cloudflare to expose a public "
f" Cloudflare tunnel: OFF ({_reason}). Studio is reachable on your "
"local network only. Pass --cloudflare to expose a public "
"Cloudflare HTTPS link."
)
else:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). There is no Cloudflare "
"public link. Raw port reachability was not verified; "
f" Cloudflare tunnel: OFF ({_reason}). There is no Cloudflare "
"public link. Raw port reachability was not verified; pass --cloudflare "
"to expose a public Cloudflare HTTPS link, or "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
@ -874,7 +878,9 @@ _cloudflare_url = None
_public_reachable = None
_cloudflare_requested = False
_cloudflare_flag = True
# Opt-in tri-state (mirrors the CLI): None = off by default, True = on,
# False = explicit --no-cloudflare. run_server overwrites it before the banner.
_cloudflare_flag = None
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
@ -1057,6 +1063,199 @@ def _cloudflare_tunnel_should_start(
return host in ("0.0.0.0", "::") and not api_only
def _stream_isatty(stream) -> bool:
"""isatty() that treats broken streams as non-interactive.
isatty() can raise under service wrappers (closed stdin -> ValueError;
sys.stdin None in Windows GUI -> AttributeError); such a stream can't host a
prompt, which is a fallback, not an error.
"""
try:
return stream.isatty()
except (AttributeError, ValueError):
return False
def _terminal_password_gate(
*,
tunnel_will_start: bool,
host: str,
secure: bool,
api_only: bool,
frontend_served: bool,
is_colab: bool = False,
) -> Tuple[bool, bool]:
"""Force a terminal password change before the public tunnel goes up.
When the tunnel is about to publish Studio and the seeded admin password was
never changed, ask for a new one (masked, confirmed) before any public URL
exists. The CLI normally does this before re-exec'ing the backend; this is
the backstop for direct `python run.py` launches and older-CLI installs.
Must run BEFORE the uvicorn socket binds: on a wildcard bind the served HTML
injects the bootstrap credential, so a pre-gate listener would hand the
default password to anyone reaching the raw port while the operator types.
Returns (proceed, drop_bootstrap_injection):
proceed False -> abort the launch (interactive refusal, or a headless
public launch nothing would protect); fail closed.
drop_bootstrap_injection True -> caller must null
app.state.bootstrap_password: the password just changed (stale), or a
public URL is about to serve the default credential and must not leak it.
Without a usable terminal the prompt is skipped: proceed if the bootstrap
deadline (armed later) will protect the launch; if even that is disabled
(api-only, timeout 0) nothing protects it, so refuse. NOT wrapped in a broad
try/except: an auth storage failure must abort rather than expose the default.
"""
if not tunnel_will_start:
return True, False
from auth import hashing as _auth_hashing
from auth import storage as _auth_storage
from auth.bootstrap_timeout import (
bootstrap_timeout_seconds,
should_arm_bootstrap_timeout,
)
from auth.terminal_prompt import (
prompt_for_password_change,
should_prompt_password_change,
)
_admin = _auth_storage.DEFAULT_ADMIN_USERNAME
# Gate can run before lifespan: seed the admin row here (idempotent).
_auth_storage.ensure_default_admin()
requires_change = _auth_storage.requires_password_change(_admin)
if not requires_change:
return True, False
if not should_prompt_password_change(
tunnel_will_start = tunnel_will_start,
requires_change = requires_change,
stdin_isatty = _stream_isatty(sys.stdin),
stderr_isatty = _stream_isatty(sys.stderr),
):
# No terminal: only proceed if the bootstrap deadline will arm; api-only
# and TIMEOUT=0 never arm it, leaving the default credential public.
deadline_arms = should_arm_bootstrap_timeout(
host = host,
secure = secure,
api_only = api_only,
frontend_served = frontend_served,
is_colab = is_colab,
requires_change = True,
timeout_seconds = bootstrap_timeout_seconds(),
)
if not deadline_arms:
print(
"Refusing to publish Studio on a public Cloudflare URL: the "
"default admin password was never changed, no terminal is "
"attached to change it here, and the bootstrap shutdown "
"deadline does not apply to this launch (api-only, or "
"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0). Change the password "
"first (run `unsloth studio` locally and log in, or re-run "
"with a terminal attached), then retry.",
file = sys.stderr,
flush = True,
)
return False, False
# The public page won't auto-fill the bootstrap credential (suppressed
# below) and the seeded file may already be gone, so point recovery at a
# terminal-attached run / reset-password instead of reading it from disk.
print(
" WARNING: the default admin password is still active while "
"Studio is about to be published on a public Cloudflare URL, and "
"no terminal is attached to change it here. The public page will "
"NOT auto-fill the bootstrap credential. Set a new password by "
"running `unsloth studio` locally with a terminal attached, or "
"`unsloth studio reset-password`. Studio shuts down after the "
"bootstrap deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT, default 1h) "
"unless the password is changed.",
file = sys.stderr,
flush = True,
)
# Never serve the default credential in HTML over a public URL.
return True, True
def _is_current_password(candidate: str) -> bool:
record = _auth_storage.get_user_and_secret(_admin)
if record is None:
return False
salt, pwd_hash, _jwt_secret, _must_change = record
return _auth_hashing.verify_password(candidate, salt, pwd_hash)
def _apply_change(new_password: str) -> None:
# Same effects as routes/auth.py change_password: rehash, rotate the JWT
# secret, revoke refresh tokens in the SAME transaction.
_auth_storage.update_password(_admin, new_password, revoke_refresh_tokens = True)
changed = prompt_for_password_change(
min_length = _auth_storage.MIN_PASSWORD_LENGTH,
is_current_password = _is_current_password,
apply_change = _apply_change,
out = sys.stderr,
)
return (True, True) if changed else (False, False)
def _apply_supplied_password(password_value: "Optional[str]") -> None:
"""Non-interactively set the INITIAL admin password before the socket binds,
for a direct ``python run.py`` launch (the CLI does this in its own parent).
Value comes from --password / UNSLOTH_STUDIO_PASSWORD / stdin.
Only ever sets the FIRST password: an already-set one is a hard error, an
invalid value fails closed. NOT wrapped in a broad try/except: an auth
storage failure must abort rather than expose the default credential.
"""
from auth import hashing as _auth_hashing
from auth import storage as _auth_storage
from auth.terminal_prompt import SUPPLIED_PASSWORD_ENV, resolve_supplied_password
supplied = resolve_supplied_password(password_value)
# Strip the env var once read so child subprocesses (cloudflared, llama-server,
# code-exec tools) can't inherit the plaintext via /proc/PID/environ. Mirrors
# the CLI. Unconditional: strips a leftover value even when a literal --password won.
os.environ.pop(SUPPLIED_PASSWORD_ENV, None)
if not supplied:
return
_admin = _auth_storage.DEFAULT_ADMIN_USERNAME
_auth_storage.ensure_default_admin()
if not _auth_storage.requires_password_change(_admin):
print(
"Error: a Studio admin password is already set; --password only sets "
"the initial password. Run `unsloth studio reset-password` first.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
def _is_current_password(candidate: str) -> bool:
record = _auth_storage.get_user_and_secret(_admin)
if record is None:
return False
salt, pwd_hash, _jwt_secret, _must_change = record
return _auth_hashing.verify_password(candidate, salt, pwd_hash)
if len(supplied) < _auth_storage.MIN_PASSWORD_LENGTH:
print(
f"Error: password must be at least {_auth_storage.MIN_PASSWORD_LENGTH} "
"characters; not starting.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
if _is_current_password(supplied):
print(
"Error: the new password must differ from the current bootstrap "
"password; not starting.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
_auth_storage.update_password(_admin, supplied, revoke_refresh_tokens = True)
print(f"Password updated for '{_admin}'.", file = sys.stderr, flush = True)
def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
"""Honor an explicit --enable-tools/--disable-tools; None leaves the policy
unset (tools default on, per-request enable_tools honored). Host is never
@ -1075,9 +1274,10 @@ def run_server(
silent: bool = False,
api_only: bool = False,
llama_parallel_slots: int = 1,
cloudflare: bool = True,
cloudflare: "Optional[bool]" = None,
secure: bool = False,
enable_tools: "Optional[bool]" = None,
password: "Optional[str]" = None,
emit_tauri_port: bool = True,
):
"""
@ -1090,6 +1290,9 @@ def run_server(
silent: Suppress startup messages
api_only: API server only, no frontend (for Tauri desktop app)
llama_parallel_slots: parallel slots for llama-server
cloudflare: opt in to the public Cloudflare HTTPS tunnel for a wildcard
bind. Tri-state: None (unset) and False both mean off; True enables it.
--secure implies it (True) and rejects an explicit False.
enable_tools: explicit --enable-tools/--disable-tools policy; None leaves
the default (tools on, per-request enable_tools honored)
emit_tauri_port: print the machine-readable TAURI_PORT line the desktop
@ -1111,13 +1314,16 @@ def run_server(
initialize_parent_lifetime()
# --secure exposes only the Cloudflare link: force a loopback bind so the raw
# port is never public (even with -H 0.0.0.0), and reject the contradictory combo.
if secure and not cloudflare:
raise SystemExit(
"A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link"
)
# --secure exposes ONLY the Cloudflare link: reject --secure --no-cloudflare,
# then force a loopback bind so the raw port is never public (even -H 0.0.0.0).
# Otherwise keep the tri-state so the banner distinguishes "off by default"
# from an explicit --no-cloudflare.
if secure:
if cloudflare is False:
raise SystemExit(
"--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare."
)
cloudflare = True
host = "127.0.0.1"
# `unsloth studio run` installs its own resolved policy and passes None here.
@ -1156,6 +1362,15 @@ def run_server(
from threading import Thread, Event
import uvicorn
# `from main import app` below loads torch/unsloth/transformers (~2 min cold,
# silent), so print a flushed heads-up (piped stdout is block-buffered).
if not silent:
print(
"Loading Unsloth Studio, please wait... (this can take a few minutes)",
flush = True,
)
print(" - loading PyTorch, Unsloth and Transformers...", flush = True)
import_started = time.perf_counter()
from main import app, setup_frontend, _IS_COLAB
@ -1164,6 +1379,8 @@ def run_server(
"Imported FastAPI app in %.1fms",
(time.perf_counter() - import_started) * 1000,
)
if not silent:
print(" - Starting server...", flush = True)
from utils.paths import ensure_studio_directories
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
@ -1191,7 +1408,7 @@ def run_server(
print("=" * 50)
if blocker:
pid, name = blocker
print(f"Port {original_port} is already in use by " f"{name} (PID {pid}).")
print(f"Port {original_port} is already in use by {name} (PID {pid}).")
else:
print(f"Port {original_port} is already in use.")
print(f"Unsloth Studio will use port {port} instead.")
@ -1304,6 +1521,44 @@ def run_server(
app.state.trigger_shutdown = _trigger_shutdown
# A supplied --password / UNSLOTH_STUDIO_PASSWORD / stdin sets the initial
# admin password before the gate and socket bind (direct `python run.py`;
# the CLI applies it in its own parent).
_apply_supplied_password(password)
# Never publish with the seeded default password active: prompt first (or
# warn / fail closed headless; see _terminal_password_gate). Runs BEFORE the
# socket binds so a pre-gate listener can't hand out the injected credential.
_pw_proceed, _pw_drop_bootstrap = _terminal_password_gate(
tunnel_will_start = _cloudflare_tunnel_should_start(
cloudflare = cloudflare,
host = host,
secure = secure,
api_only = api_only,
is_colab = _IS_COLAB,
),
host = host,
secure = secure,
api_only = api_only,
frontend_served = bool(frontend_path) and not api_only,
is_colab = _IS_COLAB,
)
if not _pw_proceed:
print(
"Not starting Studio; set a new admin password first, or launch "
"without --secure/--cloudflare.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
if _pw_drop_bootstrap:
# Password just changed (stale) or a public URL is about to serve the
# default credential: don't leak it in the HTML. Lifespan runs AFTER this
# and re-reads the bootstrap password, so the flag (not a plain None)
# makes it skip that re-read.
app.state.suppress_bootstrap_injection = True
app.state.bootstrap_password = None
# Run server in a daemon thread with explicit new_event_loop() +
# run_until_complete() (not asyncio.run) so nest_asyncio's patches don't
# interfere when Colab/IPython already runs a loop on the main thread.
@ -1373,6 +1628,7 @@ def run_server(
is_colab = _IS_COLAB,
)
_cloudflare_requested = _cloudflare_enabled
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
@ -1460,6 +1716,14 @@ def _build_arg_parser():
default = "127.0.0.1",
help = "Host to bind to (default: 127.0.0.1; use 0.0.0.0 for network/cloud access)",
)
parser.add_argument(
"--password",
default = None,
help = "Set the INITIAL admin password non-interactively (headless), only when "
"none is set yet. Also reads UNSLOTH_STUDIO_PASSWORD, or --password - for stdin. "
"A literal value is visible in the process list. Rotate later via "
"`unsloth studio reset-password`.",
)
parser.add_argument("--port", type = int, default = 8888, help = "Port to bind to")
parser.add_argument(
"--frontend",
@ -1476,11 +1740,13 @@ def _build_arg_parser():
parser.add_argument(
"--cloudflare",
action = argparse.BooleanOptionalAction,
default = True,
help = "Auto-create a free Cloudflare HTTPS tunnel for non-api-only wildcard "
"binds (0.0.0.0 or ::), exposing Studio on a PUBLIC internet URL (default on). "
"Pass --no-cloudflare to disable that Cloudflare URL; it does not change a "
"public wildcard bind. --api-only keeps it off unless paired with --secure.",
default = None,
help = "Expose Studio on a PUBLIC internet URL via a free Cloudflare HTTPS "
"tunnel, for non-api-only wildcard binds (0.0.0.0 or ::). Off by default; "
"pass --cloudflare to enable it (--secure implies it), --no-cloudflare to "
"force it off. It does not change a raw wildcard bind. If the admin "
"password was never changed, Studio asks for a new one in the terminal "
"before publishing the URL.",
)
parser.add_argument(
"--secure",
@ -1488,7 +1754,9 @@ def _build_arg_parser():
default = False,
help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
"if the tunnel can't start. Without it, --no-secure also serves the raw "
"0.0.0.0 port, which is reachable from anywhere on the network",
"0.0.0.0 port, which is reachable from anywhere on the network. If the "
"admin password was never changed, Studio asks for a new one in the "
"terminal before publishing the URL.",
)
# Back-compat: accept --not-secure as a hidden alias for --no-secure.
parser.add_argument(
@ -1550,7 +1818,7 @@ if __name__ == "__main__":
args = parser.parse_args()
if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
if args.secure and not args.cloudflare:
if args.secure and args.cloudflare is False:
parser.error(
"--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare"
)
@ -1564,6 +1832,7 @@ if __name__ == "__main__":
cloudflare = args.cloudflare,
secure = args.secure,
enable_tools = args.enable_tools,
password = args.password,
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)

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

@ -27,7 +27,7 @@ from utils.paths import (
project_workspaces_root,
studio_db_path,
)
from utils.paths.external_media import is_linux_run_media_path
from utils.paths.external_media import is_linux_run_media_path, is_local_filesystem_root
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
@ -69,6 +69,25 @@ def _denied_path_prefixes() -> list[str]:
return []
def is_denied_system_path(path: str) -> bool:
"""True if *path* is, or descends from, a denied system directory.
Mirrors the denylist add_scan_folder() enforces at registration so the
browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds
a broad root (a Windows drive root C:\\ or a legacy-registered / root). The
/run carve-out keeps Linux removable-media mounts browseable. Expects an
already-resolved (realpath) path so symlinks cannot escape into a denied subtree.
"""
is_win = platform.system() == "Windows"
check = os.path.normcase(path) if is_win else path
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
if prefix == "/run" and is_linux_run_media_path(check):
continue
return True
return False
def _contains_sensitive_path_component(path: str) -> bool:
return _shared_contains_sensitive_path_component(path)
@ -931,6 +950,12 @@ def add_scan_folder(path: str) -> dict:
raise ValueError("Path must be a directory, not a file")
if not os.access(normalized, os.R_OK | os.X_OK):
raise ValueError("Path is not readable")
# Reject a local filesystem root ("/", or a bare Windows drive root "C:\\"):
# registering one seeds the browse allowlist with a root above denied system
# dirs. A UNC share root (\\server\share) has none under it and was
# registerable before this guard, so it stays allowed. Mirrors scan_folders.py.
if is_local_filesystem_root(normalized):
raise ValueError("The filesystem root cannot be registered")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")

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

@ -1631,6 +1631,48 @@ class TestAnthropicMessagesToolRouting:
assert entry["status"] == "cancelled"
assert monitor.active_count() == 0
@staticmethod
def _sse_blob(chunks):
# StreamingResponse may hand back str or already-encoded bytes.
return "".join(c.decode() if isinstance(c, (bytes, bytearray)) else c for c in chunks)
def test_plain_streaming_unclassified_error_emits_error_event(self, monkeypatch):
# An unclassified mid-stream failure must surface as an SSE `error` event
# and stop, not a message_stop that masks a truncated turn as clean.
def _gen_boom(**_kwargs):
yield "partial"
raise RuntimeError("llama-server crashed mid-decode")
_mock_backend(monkeypatch, generate_chat_completion = _gen_boom)
payload = _basic_payload(stream = True)
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
blob = self._sse_blob(self._consume_response(response))
assert "event: error" in blob
assert '"type": "error"' in blob
assert "event: message_stop" not in blob
def test_tool_streaming_unclassified_error_emits_error_event(self, monkeypatch):
# Same guarantee on the tool-calling stream path.
def _gen_tools_boom(**_kwargs):
yield {"type": "content", "text": "partial"}
raise RuntimeError("llama-server crashed mid-decode")
_mock_backend(monkeypatch, generate_chat_completion_with_tools = _gen_tools_boom)
payload = _basic_payload(
stream = True,
enable_tools = True,
tools = [{"type": "web_search_20250305", "name": "web_search"}],
)
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
blob = self._sse_blob(self._consume_response(response))
assert "event: error" in blob
assert '"type": "error"' in blob
assert "event: message_stop" not in blob
def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch):
_mock_backend(monkeypatch)
payload = _basic_payload(
@ -1739,7 +1781,9 @@ class TestAnthropicMessagesToolRouting:
assert backend.calls[0][0] == "plain"
def test_server_tool_alias_enters_tool_path_when_policy_unset(self, monkeypatch):
# Mirror of the previous test for the default (None) policy.
# Mirror of the previous test for the default (None) policy. An omitted
# permission_mode still runs here because web_search is a safe server tool
# (only a selected terminal/python would require the missing gate).
backend = _mock_backend(monkeypatch)
payload = _basic_payload(
tools = [{"type": "web_search_20250305", "name": "web_search"}],
@ -1761,6 +1805,126 @@ class TestAnthropicMessagesToolRouting:
assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"]
assert backend.calls == []
def test_permission_mode_gating_for_server_tools(self, monkeypatch):
# ask is a request for a per-call pause this channel cannot honor, so it is
# always rejected, even for a safe-only server tool (web_search).
safe_tools = [{"type": "web_search_20250305", "name": "web_search"}]
backend = _mock_backend(monkeypatch)
payload = _basic_payload(tools = safe_tools, permission_mode = "ask")
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "no confirmation channel" in exc.value.detail["error"]["message"]
assert backend.calls == []
# auto only gates unsafe calls, so a safe-only selection runs (nothing to
# gate), like the omitted default. Both keep existing callers working.
for extra in ({"permission_mode": "auto"}, {}):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(tools = safe_tools, **extra)
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
# But auto or an omitted mode that would run a local tool (terminal/python,
# via a bare Anthropic tool type or enabled_tools) is rejected, since that
# tool could need the gate this channel lacks.
for local_payload in (
_basic_payload(tools = [{"type": "terminal", "name": "terminal"}]),
_basic_payload(
tools = [{"type": "terminal", "name": "terminal"}], permission_mode = "auto"
),
_basic_payload(tools = safe_tools, enable_tools = True, enabled_tools = ["python"]),
):
backend = _mock_backend(monkeypatch)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(local_payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "terminal" in exc.value.detail["error"]["message"]
assert backend.calls == []
# off, full, and a legacy confirm_tool_calls=False opt-out all run, even
# with a local tool selected. The explicit opt-out wins over the mode
# (mirrors _permission_mode_confirm and the GGUF path), so it runs even
# under ask, which otherwise always rejects.
for extra in (
{"tools": safe_tools, "permission_mode": "off"},
{"tools": safe_tools, "permission_mode": "full"},
{"tools": safe_tools, "enabled_tools": ["python"], "confirm_tool_calls": False},
{"tools": safe_tools, "permission_mode": "ask", "confirm_tool_calls": False},
{
"tools": [{"type": "terminal", "name": "terminal"}],
"permission_mode": "ask",
"confirm_tool_calls": False,
},
):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(**extra)
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
def test_render_html_gated_for_server_tools(self, monkeypatch):
# render_html is no longer unconditionally safe: a networked canvas prompts
# in auto and this channel cannot present that gate, so selecting it under
# ask/auto/omitted rejects like terminal/python; off/full (and an explicit
# confirm opt-out) run it.
rh = {"enable_tools": True, "enabled_tools": ["render_html"]}
for mode in ("ask", "auto", None):
backend = _mock_backend(monkeypatch)
fields = dict(rh)
if mode is not None:
fields["permission_mode"] = mode
payload = _basic_payload(**fields)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "no confirmation channel" in exc.value.detail["error"]["message"]
assert backend.calls == []
for extra in (
{"permission_mode": "off"},
{"permission_mode": "full"},
{"confirm_tool_calls": False},
):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(**{**rh, **extra})
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
def test_permission_mode_rejected_before_auto_switch(self, monkeypatch):
# The unsupported-mode rejection must run before _maybe_auto_switch_model,
# so an invalid confirm-gated request never evicts the resident model
# (mirrors the pre-switch malformed- and mixed-tool guards).
import routes.inference as inf_mod
switch_calls = []
async def _rec_switch(*_args, **_kwargs):
switch_calls.append(1)
monkeypatch.setattr(inf_mod, "_maybe_auto_switch_model", _rec_switch)
safe_tools = [{"type": "web_search_20250305", "name": "web_search"}]
local_tools = [{"type": "terminal", "name": "terminal"}]
# ask (any server tool), auto with a local tool, and an omitted mode
# selecting a local tool are all rejected up front, before the switch runs.
for payload in (
_basic_payload(tools = safe_tools, permission_mode = "ask"),
_basic_payload(tools = local_tools, permission_mode = "auto"),
_basic_payload(tools = local_tools),
):
switch_calls.clear()
_mock_backend(monkeypatch)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert switch_calls == [], "rejection must precede the auto-switch"
# A supported request (off) still reaches the switch and runs the loop.
switch_calls.clear()
_mock_backend(monkeypatch)
payload = _basic_payload(tools = safe_tools, permission_mode = "off")
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert switch_calls == [1]
def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(
@ -1770,3 +1934,523 @@ class TestAnthropicMessagesToolRouting:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "plain"
def test_resumed_session_thinking_and_null_content_do_not_400():
# A resumed session replays assistant turns with `thinking` (and sometimes null)
# content. Those must be accepted (thinking dropped by the converter), not 400ed.
from pydantic import ValidationError
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "secret reasoning", "signature": "s"},
{"type": "text", "text": "the answer"},
{"type": "tool_use", "id": "t1", "name": "f", "input": {}},
],
},
{"role": "assistant", "content": None}, # tool-only turn serialized as null
],
)
# Known blocks still parse as their typed models; only the unknown one is loose.
assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock"
assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock"
assert req.messages[2].content == "" # null coerced
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content"))
assert assistant["content"] == "the answer"
assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded
# A malformed KNOWN block still fails cleanly instead of being swallowed.
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}],
)
def test_user_null_content_rejected():
# The null->"" leniency is assistant-only; a null user content must be rejected
# at the boundary, not coerced into an empty prompt and forwarded to the model.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": None}],
)
def test_user_unknown_block_rejected_not_silently_dropped():
# The converter skips user blocks it cannot translate, so a user turn whose only
# block is unknown would validate yet forward no content. Reject at the boundary
# to avoid that silent data loss (the assistant fallback is unaffected).
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "document", "source": {}}]},
],
)
def test_user_translatable_blocks_still_accepted():
# text / image / tool_result are translatable, so a real user message built from
# them must still pass; the unknown-block guard only trips on other types.
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": "AA"},
},
{"type": "tool_result", "tool_use_id": "t1", "content": "ok"},
],
}
],
)
assert [type(b).__name__ for b in req.messages[0].content] == [
"AnthropicTextBlock",
"AnthropicImageBlock",
"AnthropicToolResultBlock",
]
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai)
def test_user_malformed_known_block_still_rejected():
# The guard only allow-lists a user block's *type*; the union still validates its
# shape, so a known-but-malformed block (tool_result without tool_use_id) fails.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
],
)
def test_user_content_block_non_string_type_rejected_cleanly():
# A user block whose `type` is a non-string (unhashable list / dict, or a stray
# int) must fail as a clean validation error, not raise TypeError from the
# frozenset membership test and escape as a 500.
from pydantic import ValidationError
for bad_type in ([], {}, 5):
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": [{"type": bad_type}]}],
)
def test_assistant_missing_content_key_still_rejected():
# The null -> "" leniency is only for an EXPLICIT null. An assistant message that
# omits content entirely stays malformed and must fail required-field validation.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant"}],
)
# An explicit null is still accepted and coerced (regression guard).
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": None},
],
)
assert req.messages[1].content == ""
def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch):
# user -> assistant(null) -> user is now accepted: the null assistant turn coerces
# to "" and is dropped. The route must then coalesce the two remaining user turns
# so a strict GGUF chat template does not 400 on non-alternating roles.
backend = _mock_backend(monkeypatch, context_length = 2048)
class _Req:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/messages")
method = "POST"
async def is_disconnected(self):
return False
payload = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "first question"},
{"role": "assistant", "content": None},
{"role": "user", "content": "please continue"},
],
)
response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t"))
assert response.status_code == 200
[(_path, kwargs)] = backend.calls
user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"]
assert len(user_turns) == 1 # the two user turns were merged, not left adjacent
merged = user_turns[0]["content"]
if isinstance(merged, list):
merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
assert "first question" in merged and "please continue" in merged
def test_disable_parallel_tool_use_forwards_heartbeats_while_dropping():
"""Heartbeats from a parallel-disabled, dropped tool call must still reach
the client as SSE keepalives: the dropped call runs server-side and the
stall keepalive never fires while the generator keeps producing events, so
swallowing them recreates the silent window keepalives exist to prevent."""
import threading as _threading
from routes.inference import (
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
_anthropic_tool_stream,
)
def run_gen():
def gen():
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_0",
"arguments": {},
}
yield {"type": "heartbeat"}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_0",
"result": "r1",
}
# Second call: dropped by disable_parallel_tool_use, still executed
# server-side (heartbeats + live output).
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_1",
"arguments": {},
}
yield {"type": "heartbeat"}
yield {
"type": "tool_output",
"tool_name": "python",
"tool_call_id": "call_1",
"text": "x",
}
yield {"type": "heartbeat"}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_1",
"result": "r2",
}
yield {"type": "content", "text": "final answer"}
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_tool_stream(
request,
_threading.Event(),
run_gen,
"msg_hb",
"m",
disable_parallel_tool_use = True,
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
# One heartbeat inside the kept call, two inside the dropped window.
assert len(keepalives) >= 3
# The dropped call must not surface as a second tool_use block.
tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c]
assert len(tool_use_starts) == 1
def test_dropped_tool_output_events_emit_rate_limited_keepalives(monkeypatch):
"""A chatty tool streaming tool_output/tool_args with no heartbeats keeps the
generator busy (stall keepalive never fires); the Anthropic path can't
translate those events and drops them. Dropping silently would let an idle
proxy kill the stream, so the drop branch emits a rate-limited keepalive."""
import threading as _threading
import routes.inference as inf_mod
from routes.inference import (
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
_anthropic_tool_stream,
)
# Deterministic clock: only the drop-branch keepalive uses time.monotonic
# here, so jumping past the stall window per call makes each dropped event
# cross the rate-limit threshold. asyncio.wait uses the loop clock and
# next(gen) returns promptly, so the outer stall keepalive never fires --
# every keepalive here is from the drop branch.
_real_time = inf_mod.time
_tick = {"v": 0.0}
def _fast_monotonic():
_tick["v"] += 100.0
return _tick["v"]
fake_time = SimpleNamespace(
monotonic = _fast_monotonic,
sleep = _real_time.sleep,
time = _real_time.time,
perf_counter = _real_time.perf_counter,
)
monkeypatch.setattr(inf_mod, "time", fake_time)
n_output = 4
def run_gen():
def gen():
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_0",
"arguments": {},
}
# Chatty streamed stdout, no heartbeats.
for i in range(n_output):
yield {
"type": "tool_output",
"tool_name": "python",
"tool_call_id": "call_0",
"text": f"line {i}\n",
}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_0",
"result": "done",
}
yield {"type": "content", "text": "final answer"}
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_tool_stream(
request,
_threading.Event(),
run_gen,
"msg_drop_ka",
"m",
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
assert len(keepalives) == n_output
# Final answer still reaches the client (drop is transport-only).
assert any("final answer" in c for c in chunks)
def test_parallel_disabled_dropped_call_output_emits_rate_limited_keepalives(monkeypatch):
"""Under disable_parallel_tool_use a chatty second call is dropped whole
(drop_until_tool_end). Its tool_output/tool_args events must still emit
rate-limited keepalives: the drop window can last minutes with no heartbeats
and no stall keepalive, so swallowing them silently would let an idle proxy
kill the stream. The keepalive branch runs before the drop skip."""
import threading as _threading
import routes.inference as inf_mod
from routes.inference import (
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
_anthropic_tool_stream,
)
# Deterministic clock: jumps past the stall window per call (see sibling test).
_real_time = inf_mod.time
_tick = {"v": 0.0}
def _fast_monotonic():
_tick["v"] += 100.0
return _tick["v"]
fake_time = SimpleNamespace(
monotonic = _fast_monotonic,
sleep = _real_time.sleep,
time = _real_time.time,
perf_counter = _real_time.perf_counter,
)
monkeypatch.setattr(inf_mod, "time", fake_time)
n_output = 4
def run_gen():
def gen():
# First (kept) call.
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_0",
"arguments": {},
}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_0",
"result": "r1",
}
# Second call: dropped whole by disable_parallel_tool_use but still
# executed server-side, streaming chatty stdout with no heartbeats.
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_1",
"arguments": {},
}
for i in range(n_output):
yield {
"type": "tool_output",
"tool_name": "python",
"tool_call_id": "call_1",
"text": f"line {i}\n",
}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_1",
"result": "r2",
}
yield {"type": "content", "text": "final answer"}
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_tool_stream(
request,
_threading.Event(),
run_gen,
"msg_drop_ka2",
"m",
disable_parallel_tool_use = True,
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
assert len(keepalives) == n_output
# The dropped call must not surface as a second tool_use block.
tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c]
assert len(tool_use_starts) == 1
assert any("final answer" in c for c in chunks)
def test_plain_stream_emits_keepalive_during_prompt_stall(monkeypatch):
"""No-tool Anthropic stream must emit SSE keepalives while a long prompt
prefill blocks next(gen), matching the tool stream (finding 5). The old
single unbounded to_thread(next, ...) could sit silent past a proxy idle cap."""
import threading as _threading
import time as _time
from routes import inference as inf_mod
from routes.inference import _OPENAI_PASSTHROUGH_SSE_KEEPALIVE, _anthropic_plain_stream
monkeypatch.setattr(inf_mod, "_LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S", 0.05)
def run_gen():
def gen():
_time.sleep(0.24) # stall past several shortened keepalive windows
yield "hello world"
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_plain_stream(
request, _threading.Event(), run_gen, "msg_plain_ka", "m"
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
assert len(keepalives) >= 2
assert any("hello world" in c for c in chunks)
def test_plain_stream_closes_generator_on_disconnect():
"""On disconnect the no-tool teardown must drain any pending worker and close
the generator (finding 6). The old finally only stopped the disconnect
watcher, leaking the generator. A fake generator records close() so the
teardown is asserted deterministically, not via GC."""
import threading as _threading
from routes.inference import _anthropic_plain_stream
closed = _threading.Event()
class _FakeGen:
def __init__(self):
self._items = iter(["tok0", "tok1", "tok2", "tok3"])
def __next__(self):
return next(self._items)
def close(self):
closed.set()
def run_gen():
return _FakeGen()
state = {"disconnected": False}
async def _drive():
async def _is_disconnected():
return state["disconnected"]
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_plain_stream(
request, _threading.Event(), run_gen, "msg_plain_close", "m"
)
out = []
async for chunk in resp.body_iterator:
out.append(chunk)
if "tok0" in chunk:
# Client drops after the first token; the next loop turn tears down.
state["disconnected"] = True
return out
asyncio.run(_drive())
assert closed.is_set()

371
studio/backend/tests/test_browse_denylist.py Normal file
View file

@ -0,0 +1,371 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""System-directory denylist enforcement for the folder browser.
Once the allowlist can hold a whole Windows drive root (C:\\) or a legacy /
root, the browse endpoints must re-apply the ``_denied_path_prefixes()`` policy
``add_scan_folder`` enforces, so /etc, /proc, C:\\Windows, C:\\Program Files stay
unbrowseable even under an allowlisted root. Windows/macOS branches run on this
POSIX host by AST-extracting the pure helper with ``ntpath`` / a mocked ``platform``.
"""
from __future__ import annotations
import ast
import ntpath
import os
import posixpath
from pathlib import Path
from types import SimpleNamespace
from typing import Optional
import pytest
from hub.storage import scan_folders
from storage import studio_db
from utils.paths.external_media import is_local_filesystem_root
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
class _HTTPException(Exception):
def __init__(self, status_code: int, detail: str):
super().__init__(detail)
self.status_code = status_code
self.detail = detail
def _extract_is_denied_windows():
"""is_denied_system_path (+ _denied_path_prefixes) from studio_db.py under Windows semantics (ntpath) on a POSIX host."""
src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8")
tree = ast.parse(src)
funcs = [
n
for n in tree.body
if isinstance(n, ast.FunctionDef)
and n.name in {"_denied_path_prefixes", "is_denied_system_path"}
]
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
win_os = SimpleNamespace(
sep = "\\",
environ = {
"SystemRoot": r"C:\Windows",
"ProgramFiles": r"C:\Program Files",
"ProgramFiles(x86)": r"C:\Program Files (x86)",
},
path = SimpleNamespace(normcase = ntpath.normcase),
)
ns = {
"os": win_os,
"platform": SimpleNamespace(system = lambda: "Windows"),
# /run has no Windows analog, so the carve-out is never reached.
"is_linux_run_media_path": lambda _p: False,
}
exec(compile(module, "<extracted studio_db.py>", "exec"), ns)
return ns["is_denied_system_path"]
# is_denied_system_path -- Linux (real helper, this host)
@pytest.mark.parametrize(
"path",
[
"/etc",
"/etc/ssl/private",
"/proc",
"/proc/1",
"/sys",
"/dev",
"/boot",
"/run",
"/run/systemd/private",
"/run/media",
"/run/media/dspofu",
],
)
def test_is_denied_system_path_linux_denies_system_dirs(monkeypatch, path):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
assert studio_db.is_denied_system_path(path) is True
@pytest.mark.parametrize(
"path",
["/run/media/dspofu/nvmeB", "/run/media/dspofu/nvmeB/models"],
)
def test_is_denied_system_path_linux_allows_run_media_mounts(monkeypatch, path):
# The /run/media/<user>/<volume> carve-out keeps removable media browseable.
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
assert studio_db.is_denied_system_path(path) is False
@pytest.mark.parametrize(
"path",
["/etc-backup", "/etcetera", "/home/u/models", "/mnt/data", "/devices", "/", "/opt/models"],
)
def test_is_denied_system_path_linux_allows_non_system(monkeypatch, path):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
assert studio_db.is_denied_system_path(path) is False
def test_legacy_and_hub_denylist_agree(monkeypatch):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux")
for p in ["/etc", "/proc/1", "/home/u", "/boot", "/opt/x"]:
assert studio_db.is_denied_system_path(p) == scan_folders.is_denied_system_path(p)
# is_denied_system_path -- Windows (ntpath-backed), case-insensitive + collisions
@pytest.mark.parametrize(
"path",
[
r"C:\Windows",
r"C:\Windows\System32",
r"c:\windows",
r"C:\WINDOWS\Temp",
r"C:\Program Files",
r"C:\Program Files\x",
r"C:\Program Files (x86)\y",
r"c:\program files",
],
)
def test_is_denied_system_path_windows_denies_system_dirs(path):
is_denied = _extract_is_denied_windows()
assert is_denied(path) is True
@pytest.mark.parametrize(
"path",
[
r"C:\Models",
r"D:\models",
r"C:\WindowsApps",
r"C:\ProgramData",
r"C:\Program Files Extra",
r"E:\gguf",
r"C:\Users\me\models",
],
)
def test_is_denied_system_path_windows_allows_non_system(path):
is_denied = _extract_is_denied_windows()
assert is_denied(path) is False
# _resolve_browse_target -- real-FS integration (legacy browser)
def _extract_resolver():
"""Extract the legacy browse resolver; its inline imports use the real storage.studio_db policy."""
src = (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")
tree = ast.parse(src)
names = {
"_is_path_inside_allowlist",
"_normalize_browse_request_path",
"_browse_relative_parts",
"_match_browse_child",
"_resolve_browse_target",
}
funcs = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
ns = {
"os": os,
"Path": Path,
"Optional": Optional,
"HTTPException": _HTTPException,
"logger": SimpleNamespace(warning = lambda *a, **k: None, debug = lambda *a, **k: None),
}
exec(compile(module, "<extracted routes/models.py>", "exec"), ns)
return ns["_resolve_browse_target"]
def test_resolve_browse_target_blocks_etc_via_root():
# Registering "/" must not make /etc browsable (Codex #3 regression guard).
resolve = _extract_resolver()
with pytest.raises(_HTTPException) as exc:
resolve("/etc", [Path("/")])
assert exc.value.status_code == 403
def test_resolve_browse_target_blocks_stale_denied_root(tmp_path, monkeypatch):
# A stale scan-folder row pointing at a denied dir is refused by the
# browse-time denylist even though it is its own allowlist root. A tmp-based
# denied prefix (+ Linux compare) keeps the assertion OS-agnostic: on macOS
# tmp lives under the already-denied /private/var, masking the message.
denied = (tmp_path / "sysfake").resolve()
denied.mkdir()
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
monkeypatch.setattr(studio_db, "_denied_path_prefixes", lambda: [str(denied)])
resolve = _extract_resolver()
with pytest.raises(_HTTPException) as exc:
resolve(str(denied), [denied])
assert exc.value.status_code == 403
assert "System directories" in exc.value.detail
def test_resolve_browse_target_allows_root_itself():
resolve = _extract_resolver()
assert resolve("/", [Path("/")]) == Path("/")
def test_resolve_browse_target_allows_legit_nested_dir(tmp_path, monkeypatch):
# Force the Linux denylist so the macOS temp location (under the denied
# /private/var) doesn't reject the tmp fixture; a normal nested dir must not be over-blocked.
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
resolve = _extract_resolver()
base = tmp_path / "allowed"
sub = base / "models" / "gguf"
sub.mkdir(parents = True)
assert resolve(str(sub), [base]) == sub.resolve()
def test_resolve_browse_target_symlink_escape_blocked(tmp_path):
resolve = _extract_resolver()
base = tmp_path / "allowed"
base.mkdir()
link = base / "escape"
try:
link.symlink_to("/etc", target_is_directory = True)
except OSError:
pytest.skip("symlinks unsupported on this host")
with pytest.raises(_HTTPException) as exc:
resolve(str(link), [base])
assert exc.value.status_code == 403
# _is_path_inside_allowlist -- bare POSIX root parity (legacy == hub)
def _extract_is_inside(rel_parts, *, os_module = os):
"""Extract a standalone _is_path_inside_allowlist (os/Path only) so both browsers' copies compare without importing their heavy modules."""
src = _BACKEND_ROOT.joinpath(*rel_parts).read_text(encoding = "utf-8")
tree = ast.parse(src)
funcs = [
n
for n in tree.body
if isinstance(n, ast.FunctionDef) and n.name == "_is_path_inside_allowlist"
]
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
ns = {"os": os_module, "Path": Path}
exec(compile(module, f"<extracted {'/'.join(rel_parts)}>", "exec"), ns)
return ns["_is_path_inside_allowlist"]
# ntpath semantics with a no-FS realpath, so UNC containment can be driven on a
# POSIX CI (the real realpath cannot resolve \\server\share off Windows).
_WIN_OS = SimpleNamespace(
sep = ntpath.sep,
path = SimpleNamespace(
realpath = lambda p: ntpath.normpath(str(p)),
normcase = ntpath.normcase,
splitdrive = ntpath.splitdrive,
dirname = ntpath.dirname,
commonpath = ntpath.commonpath,
),
)
def test_legacy_and_hub_allowlist_agree_on_posix_root():
# A bare "/" allowlist entry must authorize only "/" itself in BOTH
# browsers, never descend into /var, /root, /home (which the denylist does
# not cover). Guards the hub browser against authorizing every absolute path.
legacy = _extract_is_inside(["routes", "models.py"])
hub = _extract_is_inside(["hub", "services", "models", "folder_browser.py"])
roots = [Path("/")]
for tgt in ["/var", "/root", "/home", "/usr", "/opt", "/etc"]:
assert legacy(Path(tgt), roots) is False
assert hub(Path(tgt), roots) is False
# "/" itself stays browseable; only its descendants are withheld.
assert legacy(Path("/"), roots) is True
assert hub(Path("/"), roots) is True
def test_hub_allowlist_authorizes_normal_nested_dir(tmp_path):
# The bare-root special case must not over-block a normal allowlist root's descendants.
hub = _extract_is_inside(["hub", "services", "models", "folder_browser.py"])
base = tmp_path / "allowed"
sub = base / "models" / "gguf"
sub.mkdir(parents = True)
assert hub(sub, [base]) is True
assert hub(base, [base]) is True
# add_scan_folder -- filesystem-root rejection parity (legacy == hub)
def test_legacy_add_scan_folder_rejects_filesystem_root(monkeypatch):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
with pytest.raises(ValueError, match = "filesystem root"):
studio_db.add_scan_folder("/")
def test_hub_add_scan_folder_rejects_filesystem_root(monkeypatch):
monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux")
with pytest.raises(ValueError, match = "filesystem root"):
scan_folders.add_scan_folder("/")
# is_local_filesystem_root: reject "/" and "C:\\" (roots above denied system dirs),
# but NOT a UNC share root -- registering \\server\share was allowed before this
# guard and has no system dirs under it. _pathmod drives Windows semantics on POSIX CI.
@pytest.mark.parametrize(
"path, pathmod, expected",
[
# Local filesystem roots -> rejected (True).
("/", posixpath, True),
("C:\\", ntpath, True),
("c:\\", ntpath, True),
("D:\\", ntpath, True),
# UNC share roots -> NOT a local root, stay registerable (False).
(r"\\server\share", ntpath, False),
(r"\\nas\models", ntpath, False),
("//server/share", ntpath, False),
# Device / extended-length volume roots -> still local roots (rejected),
# so neither \\?\C:\ nor a drive-letter-less \\?\Volume{GUID}\ can slip
# past the guard as if it were a share root.
(r"\\?\C:" + "\\", ntpath, True),
(r"\\.\C:" + "\\", ntpath, True),
(r"\\?\C:", ntpath, True),
(r"\\.\C:", ntpath, True),
(r"\\?\Volume{2f8e6d31-0000-0000-0000-100000000000}" + "\\", ntpath, True),
(r"\\.\Volume{2f8e6d31-0000-0000-0000-100000000000}", ntpath, True),
# Device-namespace UNC share root -> stays registerable (False).
(r"\\?\UNC\server\share", ntpath, False),
# Non-root paths (incl. deep device / extended-length) -> not a root (False).
("C:\\Models", ntpath, False),
(r"\\server\share\models", ntpath, False),
(r"\\?\C:\Users\me\models", ntpath, False),
(r"\\?\Volume{2f8e6d31-0000-0000-0000-100000000000}\models", ntpath, False),
("/home/user", posixpath, False),
],
)
def test_is_local_filesystem_root(path, pathmod, expected):
assert is_local_filesystem_root(path, _pathmod = pathmod) is expected
def test_both_guards_use_the_shared_local_root_helper():
# Register-root parity: both browsers reject the same roots via one helper, so a
# UNC-share exemption can never drift between the legacy and hub code paths.
legacy_src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8")
hub_src = (_BACKEND_ROOT / "hub" / "storage" / "scan_folders.py").read_text(encoding = "utf-8")
assert "is_local_filesystem_root(normalized)" in legacy_src
assert "is_local_filesystem_root(normalized)" in hub_src
# A registered UNC share root must authorize its own descendants in both browsers.
# os.path.commonpath raises "can't mix absolute and relative" on a bare
# \\server\share, so containment falls back to a boundary-safe prefix test; without
# it, registering a UNC share (now allowed) would 403 every folder under it.
@pytest.mark.parametrize(
"rel_parts",
[
["routes", "models.py"],
["hub", "services", "models", "folder_browser.py"],
],
)
def test_unc_share_root_authorizes_its_descendants(rel_parts):
is_inside = _extract_is_inside(rel_parts, os_module = _WIN_OS)
root = [Path(r"\\server\share")]
assert is_inside(Path(r"\\server\share"), root) is True # the root itself
assert is_inside(Path(r"\\server\share\models"), root) is True # direct child
assert is_inside(Path(r"\\server\share\a\b\c"), root) is True # deep descendant
assert is_inside(Path(r"\\SERVER\SHARE\Models"), root) is True # case-insensitive
assert is_inside(Path(r"\\server\share2\models"), root) is False # sibling share
assert is_inside(Path(r"C:\models"), root) is False # different volume

12
studio/backend/tests/test_browse_folders_route.py
View file

@ -22,6 +22,18 @@ if "structlog" not in sys.modules:
)
import routes.models as models_route
import storage.studio_db as studio_db
@pytest.fixture(autouse = True)
def _denylist_inert(monkeypatch):
# These tests exercise allowlist containment and the file-vs-directory guard,
# not the system-directory denylist (which has its own suite in
# test_browse_denylist.py). On macOS tmp_path resolves under /private/var, a
# denied prefix, so _resolve_browse_target would 403 the fixture dirs before
# the containment logic runs. Keep the denylist inert here so these
# assertions hold on every platform.
monkeypatch.setattr(studio_db, "is_denied_system_path", lambda _p: False)
def test_resolve_browse_target_returns_allowed_directory(tmp_path):

22
studio/backend/tests/test_bypass_permissions.py
View file

@ -13,6 +13,7 @@ never gating under bypass.
Run with: ``PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_bypass_permissions.py -q``
"""
import io
import os
import sys
@ -94,10 +95,23 @@ def test_safe_env_excludes_host_and_secret(monkeypatch, tmp_path):
class _FakeProc:
returncode = 0
"""A subprocess.Popen double for the drain path (``tools._drain_process_output``):
a readable ``stdout`` pipe yielding the fake output then EOF, plus
``wait()`` / ``poll()`` / ``pid``. The pid is non-existent so
``_capture_process_group``'s ``os.getpgid`` returns None; ``wait`` returns
immediately so the drain never kills.
"""
def communicate(self, timeout = None):
return ("FAKEOUT", None)
returncode = 0
# Unlikely-to-exist pid: os.getpgid raises ProcessLookupError (caught) -> None.
pid = 2**22
def __init__(self):
# Readable stdout: iter(readline, "") yields "FAKEOUT" then hits EOF.
self.stdout = io.StringIO("FAKEOUT")
def wait(self, timeout = None):
return 0
def poll(self):
return 0
@ -257,6 +271,7 @@ def test_loop_forwards_disable_sandbox_and_does_not_gate():
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):
@ -290,6 +305,7 @@ def test_loop_bypass_overrides_confirm_for_direct_callers():
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):

10
studio/backend/tests/test_chat_load_during_training.py
View file

@ -651,11 +651,19 @@ class TestLoadModelGuardIntegration(unittest.TestCase):
inf._shutdown_subprocess = MagicMock()
llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None)
llama.unload_model = MagicMock()
cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None)
cfg = SimpleNamespace(
is_gguf = False,
is_lora = False,
path = None,
base_model = None,
identifier = "unsloth/Qwen3-1.7B",
)
request = LoadRequest(model_path = "unsloth/Qwen3-1.7B")
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"}
with (
# Pin the latest-sidecar tier check so the guard path stays offline.
patch("utils.transformers_version.latest_tier_active_for", return_value = False),
patch.object(self.route, "validate_extra_args", return_value = None),
patch.object(
self.route,

34
studio/backend/tests/test_cloudflare_tunnel.py
View file

@ -691,13 +691,14 @@ def _argparse_default(source, option):
return None
def test_run_server_cloudflare_default_true():
def test_run_server_cloudflare_default_off():
defaults = _func_param_defaults(_RUN_PY.read_text(), "run_server")
assert defaults.get("cloudflare") is True
assert "cloudflare" in defaults
assert defaults["cloudflare"] is None
def test_argparse_cloudflare_default_true():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True
def test_argparse_cloudflare_default_off():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is None
def test_verify_global_reachability_marks_private_address_unreachable():
@ -832,6 +833,31 @@ def test_cloudflare_line_states_disabled_when_off(monkeypatch):
assert "local network only" in out
def test_cloudflare_line_labels_unset_as_default(monkeypatch):
# None = off by default (no flag) -> banner says "(default)", not "(--no-cloudflare)".
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = False,
cloudflare_requested = False,
cloudflare_flag = None,
)
assert "Cloudflare tunnel: OFF (default)" in out
assert "--no-cloudflare" not in out
def test_cloudflare_line_labels_explicit_no_cloudflare(monkeypatch):
# False = explicit --no-cloudflare -> banner says "(--no-cloudflare)".
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = False,
cloudflare_requested = False,
cloudflare_flag = False,
)
assert "Cloudflare tunnel: OFF (--no-cloudflare)" in out
def test_cloudflare_line_states_failed_when_requested_but_no_url(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,

50
studio/backend/tests/test_coding_agents.py Normal file
View file

@ -0,0 +1,50 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for coding-agent CLI detection used by the API-keys settings panel."""
from unittest.mock import patch
from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents
def test_matches_unsloth_start_subcommands():
# Each entry must be an actual `unsloth start <agent>` subcommand name
# (unsloth_cli/commands/start.py). Spelled out here rather than imported
# from that module, which pulls in the CLI's heavier dependencies.
assert CODING_AGENTS == ("claude", "codex", "openclaw", "opencode", "hermes", "pi")
def test_detects_only_agents_present_on_path():
installed = {"claude", "opencode"}
with patch(
"utils.coding_agents.shutil.which",
side_effect = lambda name: f"/usr/bin/{name}" if name in installed else None,
):
assert detect_installed_coding_agents() == ["claude", "opencode"]
def test_returns_empty_list_when_nothing_is_installed():
with patch("utils.coding_agents.shutil.which", return_value = None):
assert detect_installed_coding_agents() == []
def test_preserves_declared_order_regardless_of_path_lookup_order():
with patch(
"utils.coding_agents.shutil.which",
side_effect = lambda name: name if name in ("pi", "claude", "hermes") else None,
):
assert detect_installed_coding_agents() == ["claude", "hermes", "pi"]
def test_treats_a_path_lookup_error_as_not_installed():
# An advisory check: shutil.which raising for one entry (e.g. a permission
# error walking a PATH directory) should not take down the whole endpoint,
# and should not stop the remaining agents from being checked.
def flaky_which(name: str):
if name == "codex":
raise OSError("permission denied")
return name if name == "claude" else None
with patch("utils.coding_agents.shutil.which", side_effect = flaky_which):
assert detect_installed_coding_agents() == ["claude"]

314
studio/backend/tests/test_completion_masking.py Normal file
View file

@ -0,0 +1,314 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Completion-only masking policy: auto-detect first, manual table fallback.
Covers utils.datasets.completion_masking.apply_completion_masking, shared by
the CUDA trainer (core/training/trainer.py) and the MLX worker
(core/training/worker.py):
- unmapped models use chat template auto-detection (previously masking was
silently disabled),
- gpt-oss goes auto-first too (its quantized checkpoints ship a template
the manual markers cannot match),
- an auto-detection failure falls back to the template table markers,
- a table miss after an auto failure warns and leaves the trainer unchanged.
"""
from __future__ import annotations
import pytest
from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers
from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER
class _Trainer:
"""Sentinel trainer; train_fn wraps it in a new object when applied."""
class _Recorder:
"""Fake train_on_responses_only that records calls."""
def __init__(self):
self.calls = []
def __call__(self, trainer, **kwargs):
self.calls.append(kwargs)
wrapped = _Trainer()
wrapped.wrapped_from = trainer
return wrapped
def _detect_ok(processor):
return "<INS>", "<RES>"
def _detect_fail(processor):
raise ValueError(
"Unsloth: Could not reliably auto-detect response_part - "
"pass instruction_part and response_part."
)
_AUTO = {"instruction_part": "<INS>", "response_part": "<RES>"}
class _Notes:
def __init__(self):
self.messages = []
def __call__(self, level, message):
self.messages.append((level, message))
def warnings(self):
return [m for level, m in self.messages if level == "warning"]
def test_unmapped_model_uses_auto_detection():
# Unmapped model: the auto path applies masking (was silently disabled).
trainer = _Trainer()
train_fn = _Recorder()
notes = _Notes()
result, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok
)
assert applied is True
assert result.wrapped_from is trainer
assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers
assert notes.warnings() == []
def test_mapped_model_prefers_auto_detection():
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok
)
assert applied is True
assert train_fn.calls == [dict(_AUTO)]
def test_gpt_oss_uses_auto_detection_first():
# The quantized gpt-oss checkpoints ship a template without the
# <|channel|>final header, where the manual markers match nothing; auto
# derives markers from the template the checkpoint actually ships.
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok
)
assert applied is True
assert train_fn.calls == [dict(_AUTO)]
def test_gpt_oss_detection_failure_falls_back_to_manual_markers():
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail
)
assert applied is True
expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"]
assert train_fn.calls == [
{
"instruction_part": expected["instruction"],
"response_part": expected["response"],
}
]
def test_auto_failure_falls_back_to_template_table():
trainer = _Trainer()
train_fn = _Recorder()
notes = _Notes()
result, applied = apply_completion_masking(
trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail
)
assert applied is True
assert result.wrapped_from is trainer
expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]
assert train_fn.calls == [
{
"instruction_part": expected["instruction"],
"response_part": expected["response"],
},
]
assert any("falling back to the template table" in m for m in notes.warnings())
def test_application_failure_propagates_not_fallback():
# Detection succeeds; a failure while APPLYING the masking must propagate,
# never silently fall back to full-sequence training.
def train_fn(trainer, **kwargs):
raise RuntimeError("dataset map worker crashed")
with pytest.raises(RuntimeError, match = "dataset map worker crashed"):
apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok)
def test_preset_tokenizer_markers_used_directly():
# Preset unsloth marker attrs skip detection; zoo reuses them on a bare call.
class _Tok:
_unsloth_input_part = "<I>"
_unsloth_output_part = "<O>"
trainer = _Trainer()
trainer.processing_class = _Tok()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail
)
assert applied is True
assert train_fn.calls == [{}] # bare call, stored parts
def test_table_miss_warns_and_disables_without_crashing():
trainer = _Trainer()
train_fn = _Recorder()
notes = _Notes()
result, applied = apply_completion_masking(
trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail
)
assert applied is False
assert result is trainer # unchanged: full sequence training
assert train_fn.calls == [] # detection failed; nothing applied
assert any("could not be applied" in m for m in notes.warnings())
assert any("full sequences" in m for m in notes.warnings())
def test_num_proc_forwarded_only_when_given():
# CUDA path passes num_proc; the MLX path omits it.
train_fn = _Recorder()
apply_completion_masking(
_Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok
)
assert train_fn.calls == [dict(_AUTO, num_proc = 4)]
train_fn = _Recorder()
apply_completion_masking(
_Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail
)
assert train_fn.calls[0]["num_proc"] == 4
train_fn = _Recorder()
apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok)
assert train_fn.calls == [dict(_AUTO)]
def test_manual_fallback_failure_propagates_to_caller():
# Errors while applying the manual fallback must propagate to the caller.
def train_fn(trainer, **kwargs):
raise RuntimeError("boom")
with pytest.raises(RuntimeError, match = "boom"):
apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn)
def test_notify_is_optional():
train_fn = _Recorder()
_, applied = apply_completion_masking(
_Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail
)
assert applied is False
def test_lookup_manual_markers():
template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B")
assert template == "qwen3"
assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"]
assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"]
template, instruction, response = lookup_manual_markers("some-org/unknown")
assert (template, instruction, response) == (None, None, None)
template, instruction, response = lookup_manual_markers(None)
assert (template, instruction, response) == (None, None, None)
def test_renamed_gpt_oss_gets_template_markers():
# Name-detected as gpt-oss but not in the exact-name table: must use the
# gpt-oss markers, not fall through to full-sequence training.
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail
)
assert applied is True
expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"]
assert train_fn.calls == [
{
"instruction_part": expected["instruction"],
"response_part": expected["response"],
}
]
class _FakeTokenizerWrapper:
"""mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped
tokenizer, underscore attrs do not (so preset markers are hidden)."""
def __init__(self, tokenizer):
object.__setattr__(self, "_tokenizer", tokenizer)
def __getattr__(self, attr):
if attr.startswith("_"):
return object.__getattribute__(self, attr)
return getattr(object.__getattribute__(self, "_tokenizer"), attr)
_FakeTokenizerWrapper.__name__ = "TokenizerWrapper"
def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers():
# Markers live on the inner HF tokenizer that the wrapper hides; the helper
# must unwrap so the preset bare-call path still fires on MLX.
class _Tok:
_unsloth_input_part = "<I>"
_unsloth_output_part = "<O>"
trainer = _Trainer()
trainer.tokenizer = _FakeTokenizerWrapper(_Tok())
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail
)
assert applied is True
assert train_fn.calls == [{}] # bare call, stored parts
def test_mlx_tokenizer_wrapper_unwrapped_for_detection():
# Detection must see the real tokenizer, not the wrapper, so it does not
# depend on the loader's __call__ patch.
class _Tok:
pass
inner = _Tok()
trainer = _Trainer()
trainer.tokenizer = _FakeTokenizerWrapper(inner)
train_fn = _Recorder()
seen = []
def detect(processor):
seen.append(processor)
return "<INS>", "<RES>"
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect
)
assert applied is True
assert seen == [inner]

97
studio/backend/tests/test_data_recipe_seed.py
View file

@ -124,3 +124,100 @@ def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, e
assert result.status == "error"
assert result.error == "Text extraction failed."
assert _block_files(seed_route) == []
_TEST_UPLOAD_UID = "0f" * 16
def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
_run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID)
assert _block_files(seed_route, _TEST_UPLOAD_UID) != []
result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert result == {"status": "ok", "deleted": True}
assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists()
def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert result == {"status": "ok", "deleted": False}
def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block("../escape"))
assert exc.value.status_code == 400
def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
_run_upload(seed_route, "notes.txt", b"hello", block_id = "n1")
assert _block_files(seed_route, "n1") != []
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block("n1"))
assert exc.value.status_code == 400
assert _block_files(seed_route, "n1") != []
def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
outside = tmp_path / "outside"
outside.mkdir()
(outside / "victim.txt").write_text("keep me")
root = seed_route.UNSTRUCTURED_UPLOAD_ROOT
root.mkdir(parents = True)
(root / _TEST_UPLOAD_UID).symlink_to(outside)
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert exc.value.status_code == 400
assert (outside / "victim.txt").exists()
def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
root = seed_route.UNSTRUCTURED_UPLOAD_ROOT
block_dir = root / _TEST_UPLOAD_UID
block_dir.mkdir(parents = True)
(block_dir / "victim.txt").write_text("keep me")
calls = []
def noop_rmtree(path, *args, **kwargs):
calls.append((path, args, kwargs))
monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree)
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert calls
assert exc.value.status_code == 500
assert block_dir.exists()
def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10)
first = _run_upload(seed_route, "a.txt", b"123456789")
assert first.status == "ok"
with pytest.raises(seed_route.HTTPException) as exc:
_run_upload(seed_route, "b.txt", b"123")
assert exc.value.status_code == 413
# Another block starts with its own untouched budget.
other = _run_upload(seed_route, "c.txt", b"123", block_id = "other")
assert other.status == "ok"

42
studio/backend/tests/test_hf_xet_fallback.py
View file

@ -287,7 +287,9 @@ def test_degrades_when_shared_helper_import_raises_importerror():
def test_retries_under_light_gpu_init_when_import_fails(monkeypatch):
"""GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim
retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails."""
retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.
The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly
before asserting the retry/degrade behavior."""
import importlib
import os
@ -321,11 +323,15 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch):
sys.meta_path.insert(0, finder)
try:
degraded = importlib.import_module("utils.hf_xet_fallback")
# First attempt without the light env, then a retry with it set.
# Import is light (lazy backend); unsloth_zoo not loaded yet.
assert seen_env == [], seen_env
# First use of a heavy helper triggers the load (attempt without the light env, then a retry
# with it set); accessing DownloadStallError drives it via __getattr__.
stall_error = degraded.DownloadStallError
assert seen_env == [None, "1"], seen_env
# Both attempts raised -> Studio still boots in degraded mode.
assert issubclass(degraded.DownloadStallError, RuntimeError)
# The env override must not leak past the import.
assert issubclass(stall_error, RuntimeError)
# The env override must not leak past the load.
assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None
finally:
sys.meta_path.remove(finder)
@ -333,3 +339,31 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch):
sys.modules.update(saved)
if saved_shim is not None:
sys.modules["utils.hf_xet_fallback"] = saved_shim
def test_importing_child_should_disable_xet_stays_light(monkeypatch):
"""Regression guard for the stale-transformers-sidecar bug: importing the shim (and
``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls
this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here
would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend
and breaking 5.x models (Qwen3.5/GLM/gemma-4)."""
import importlib
for name in [
m
for m in list(sys.modules)
if m == "transformers"
or m.startswith("transformers.")
or m == "unsloth_zoo"
or m.startswith("unsloth_zoo.")
or m == "utils.hf_xet_fallback"
]:
monkeypatch.delitem(sys.modules, name, raising = False)
mod = importlib.import_module("utils.hf_xet_fallback")
# The lightweight decision works without the heavy backend.
assert mod.child_should_disable_xet({"disable_xet": True}) is True
assert mod.child_should_disable_xet({}) is False
# And nothing heavy was imported as a side effect.
assert "transformers" not in sys.modules, "importing the shim must not import transformers"
assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo"

123
studio/backend/tests/test_index_bootstrap_loopback.py Normal file
View file

@ -0,0 +1,123 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression coverage for bootstrap password exposure to remote clients."""
from types import SimpleNamespace
def _request(
client_host,
request_host = "127.0.0.1",
headers = None,
):
"""Build a minimal request; ``None`` models an unresolved peer / absent Host."""
client = None if client_host is None else SimpleNamespace(host = client_host, port = 0)
hdrs = {}
if request_host is not None:
hdrs["host"] = request_host
hdrs.update(headers or {})
return SimpleNamespace(client = client, headers = hdrs, url = SimpleNamespace(hostname = request_host))
def test_loopback_peers_are_local():
from main import _is_local_bootstrap_request
cases = (
("127.0.0.1", "127.0.0.1"),
("::1", "::1"),
("::ffff:127.0.0.1", "::ffff:127.0.0.1"),
("127.0.0.1", "localhost"),
)
for peer, host in cases:
assert _is_local_bootstrap_request(_request(peer, host)) is True, (peer, host)
def test_non_loopback_peers_are_remote():
from main import _is_local_bootstrap_request
# ::1%eth0 is a scope-id'd address, which ipaddress treats as loopback on
# 3.9+; it must not count as a direct local peer.
for host in ("192.168.1.10", "::ffff:192.168.1.10", "::1%eth0"):
assert _is_local_bootstrap_request(_request(host)) is False, host
def test_absent_or_unparseable_peer_fails_safe():
from main import _is_local_bootstrap_request
for host in (None, "localhost"):
assert _is_local_bootstrap_request(_request(host)) is False, host
def test_cloudflare_tunnel_clients_are_remote_despite_loopback_peer():
from main import _is_local_bootstrap_request
for client_ip in ("203.0.113.7", ""):
request = _request("127.0.0.1", headers = {"cf-connecting-ip": client_ip})
assert _is_local_bootstrap_request(request) is False, client_ip
def test_dns_rebinding_host_is_remote_despite_loopback_peer():
from main import _is_local_bootstrap_request
for host in ("attacker.example", "192.168.1.10", None):
assert _is_local_bootstrap_request(_request("127.0.0.1", host)) is False, host
def test_unparseable_request_host_fails_safe():
"""A Host that makes ``request.url.hostname`` raise must fall to remote."""
from main import _is_local_bootstrap_request
class _RaisingURL:
@property
def hostname(self):
raise ValueError("malformed host")
request = SimpleNamespace(
client = SimpleNamespace(host = "127.0.0.1", port = 0), headers = {}, url = _RaisingURL()
)
assert _is_local_bootstrap_request(request) is False
def test_reverse_proxy_forwarded_headers_are_remote():
"""A loopback proxy relaying a remote client (non-Cloudflare headers) is remote."""
from main import _is_local_bootstrap_request
for header in ("forwarded", "x-forwarded-for", "x-forwarded-host", "x-real-ip"):
request = _request("127.0.0.1", "localhost", headers = {header: "203.0.113.7"})
assert _is_local_bootstrap_request(request) is False, header
def test_malformed_or_absent_host_is_remote():
"""A malformed/absent/scope-id Host must not fall back to the loopback server address."""
from main import _is_local_bootstrap_request
# incl. bracket smuggling: [::1]evil / unclosed [::1 must not reduce to ::1
for host in (
"e_vil",
"[malformed",
"",
None,
"[::1%25eth0]:8888",
"[::1]attacker",
"[::1]evil.com",
"[::1",
"[::1]x",
):
assert _is_local_bootstrap_request(_request("127.0.0.1", host)) is False, host
def test_colab_allows_notebook_proxy_but_not_shareable_tunnel(monkeypatch):
"""Colab autofills its single-user proxy, but not a public Cloudflare link."""
import main
monkeypatch.setattr(main, "_IS_COLAB", True)
# In-notebook proxy: same-origin, no tunnel header, injects off-loopback too.
assert main._should_inject_bootstrap(_request("10.0.0.2", "colab.proxy")) is True
# Shareable Cloudflare link marks visitors with cf-connecting-ip; withhold.
tunnel = _request("127.0.0.1", "localhost", headers = {"cf-connecting-ip": "203.0.113.7"})
assert main._should_inject_bootstrap(tunnel) is False
def test_non_colab_gate_requires_local_client(monkeypatch):
"""Outside Colab the gate injects only for a direct loopback client."""
import main
monkeypatch.setattr(main, "_IS_COLAB", False)
assert main._should_inject_bootstrap(_request("127.0.0.1", "localhost")) is True
assert main._should_inject_bootstrap(_request("192.168.1.10", "localhost")) is False

2
studio/backend/tests/test_inference_dispatcher_resilience.py
View file

@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
continue
calls.append(node)
assert len(calls) == 4
assert len(calls) == 5
for call in calls:
assert any(
kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False

503
studio/backend/tests/test_install_resolve_prebuilt.py
View file

@ -1,7 +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
"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode.
"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork
by default; --published-repo overrides).
These back the in-app update for source-build (markerless) installs: the backend
asks the installer whether an official prebuilt exists for this host without
@ -24,9 +25,7 @@ if str(_studio) not in sys.path:
ilp = importlib.import_module("install_llama_prebuilt")
if not hasattr(ilp, "published_repo_for_host") or not hasattr(
ilp, "resolve_simple_install_release_plans"
):
if not hasattr(ilp, "resolve_simple_install_release_plans"):
pytest.skip("PR symbols not present - check branch", allow_module_level = True)
FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp
@ -56,71 +55,25 @@ def _host(**kw):
return ilp.HostInfo(**base)
def test_published_repo_for_host():
# CPU-only Linux (x64 and arm64) -> ggml-org upstream.
assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM
assert (
ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64"))
== UPSTREAM
)
# GPU Linux -> fork.
assert (
ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True))
== FORK
)
assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK
# CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle).
assert (
ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True))
== UPSTREAM
)
# GPU Windows -> fork.
assert (
ilp.published_repo_for_host(
_host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True)
)
== FORK
)
# macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS).
assert (
ilp.published_repo_for_host(
_host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64")
)
== FORK
)
# Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling).
assert (
ilp.published_repo_for_host(
_host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True
)
== FORK
)
# The tooling hint is Linux-only: Windows CPU stays on ggml-org.
assert (
ilp.published_repo_for_host(
_host(system = "Windows", is_windows = True, is_x86_64 = True),
linux_amd_tooling_present = True,
)
== UPSTREAM
)
def test_macos_intel_and_arm_both_route_to_fork():
# macOS uses the unslothai fork's own Mac prebuilts for BOTH arm64 and Intel;
# there is no longer any upstream-on-macOS default path, so the obsolete
# pre-macOS-26 pin (b9415) is gone.
assert (
ilp.published_repo_for_host(
_host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64")
)
== FORK
)
assert (
ilp.published_repo_for_host(
_host(system = "Darwin", is_macos = True, is_x86_64 = True, machine = "x86_64")
)
== FORK
def test_force_cpu_clears_all_gpu_attributes_including_intel():
# --cpu-fallback is the "select the CPU prebuilt even when a GPU is present"
# escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or
# the planner still prepends the Vulkan asset on an Intel-GPU host.
host = _host(
is_linux = True,
is_x86_64 = True,
has_usable_nvidia = True,
has_physical_nvidia = True,
has_rocm = True,
rocm_gfx_target = "gfx1100",
has_intel_gpu = True,
)
forced = ilp._apply_host_overrides(host, force_cpu = True)
assert forced.has_usable_nvidia is False
assert forced.has_physical_nvidia is False
assert forced.has_rocm is False
assert forced.rocm_gfx_target is None
assert forced.has_intel_gpu is False
def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
@ -188,15 +141,13 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys):
assert out["repo"] == FORK
def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys):
# CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the
# fork so a HIP source build is not offered an upstream CPU prebuilt.
monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True))
monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo")
def _run_resolve_capture_host(monkeypatch, capsys):
"""Drive --resolve-prebuilt and return the host the resolver was handed."""
seen = {}
def _resolver(tag, host, repo, published_release_tag):
seen["repo"] = repo
seen["host"] = host
raise ilp.PrebuiltFallback("no asset")
monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
@ -207,10 +158,33 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys):
)
assert ilp.main() == ilp.EXIT_SUCCESS
out = json.loads(capsys.readouterr().out.strip().splitlines()[-1])
return seen, out
def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys):
# CPU-only Linux host (no GPU): the dispatch routes to the fork, which now
# ships the CPU prebuilt -- it no longer falls back to ggml-org upstream.
monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True))
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == FORK
assert out["repo"] == FORK
def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys):
# A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so
# detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe
# must NOT reclassify it as ROCm from tool presence alone and suppress the CPU
# bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source
# build. The host is left CPU-only and resolves against the fork.
monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True))
monkeypatch.setattr(
ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == FORK
assert seen["host"].has_rocm is False
# Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer
# sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection.
@ -360,3 +334,386 @@ def test_sm103_host_drops_cuda128_windows_build():
)
kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129])
assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name]
def _upstream_release(tag, asset_names):
return {
"tag_name": tag,
"assets": [
{"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names
],
}
def test_direct_upstream_arm64_intel_prefers_vulkan():
# Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU
# second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset).
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-arm64" in kinds
assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only():
# A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical
# True, usable False) + an Intel iGPU must NOT get the Vulkan archive even
# when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES
# and could grab the reserved card. It falls through to the CPU asset.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
def test_direct_upstream_arm64_without_intel_is_cpu_only():
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64")
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-arm64"]
def test_direct_upstream_x86_intel_prefers_vulkan():
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-cpu" in kinds
def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
# The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
# libs so a valid Vulkan install is not re-flagged unhealthy every check.
choice = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz",
url = "https://example/x",
source_label = "upstream",
install_kind = "linux-vulkan",
)
groups = ilp.runtime_payload_health_groups(choice)
assert ["libggml-cpu*.so*"] in groups
assert ["libggml-cpu-*.so*"] not in groups
def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
assert repo == UPSTREAM
assert tag == ""
assert routed.has_intel_gpu is True
def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
# A pin set WITH an explicit upstream repo is already on upstream -> kept.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
_routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
assert repo == UPSTREAM
assert tag == "b9596"
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed is host
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
# Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
# An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_non_intel_unchanged():
host = _host(is_linux = True, is_x86_64 = True)
routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
assert routed is host
def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
# The --resolve-prebuilt probe must agree with the install path: an
# auto-detected Intel host resolves against upstream (Vulkan), not the fork.
monkeypatch.setattr(
ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == UPSTREAM
assert out["repo"] == UPSTREAM
# ---------------------------------------------------------------------------
# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake
# winreg module stands in for the real registry so the walk runs anywhere.
# ---------------------------------------------------------------------------
class _FakeRegKey:
def __init__(
self,
subkeys = None,
values = None,
denied = False,
):
self.subkeys = subkeys or {}
self.values = values or {}
self.denied = denied
def __enter__(self):
return self
def __exit__(self, *exc):
return False
class _FakeWinreg:
HKEY_LOCAL_MACHINE = object()
def __init__(self, root_key):
self._root_key = root_key
def OpenKey(self, parent, name):
if parent is self.HKEY_LOCAL_MACHINE:
# Pin the production constant: a typo'd class GUID must fail here,
# not silently return the fake tree.
if name != ilp._WINDOWS_DISPLAY_CLASS_KEY:
raise FileNotFoundError(name)
if self._root_key is None:
raise FileNotFoundError(name)
return self._root_key
key = parent.subkeys.get(name)
if key is None:
# Real winreg raises OSError, never KeyError, for a missing key.
raise FileNotFoundError(name)
if key.denied:
raise PermissionError(name)
return key
def QueryInfoKey(self, key):
return (len(key.subkeys), len(key.values), 0)
def EnumKey(self, key, index):
return list(key.subkeys)[index]
def QueryValueEx(self, key, value_name):
if value_name not in key.values:
raise FileNotFoundError(value_name)
return (key.values[value_name], 1)
def _probe_with_display_class(monkeypatch, adapters):
# The helper lazily does `import winreg`; plant the fake in sys.modules the
# same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path.
monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters)))
return ilp.windows_intel_gpu_in_registry()
def test_windows_intel_registry_matches_vendor_id(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678",
"DriverDesc": "Intel(R) Arc(TM) A770 Graphics",
}
),
},
)
is True
)
def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}),
},
)
is True
)
def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684",
"DriverDesc": "NVIDIA GeForce RTX 4090",
}
),
"0001": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_1002&DEV_744C",
"DriverDesc": "AMD Radeon RX 7900 XTX",
}
),
},
)
is False
)
def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch):
# The real class key carries an ACL-restricted "Properties" subkey and can
# deny access to individual adapter keys; neither may abort the walk.
assert (
_probe_with_display_class(
monkeypatch,
{
"Properties": _FakeRegKey(denied = True),
"0000": _FakeRegKey(denied = True),
"0001": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0",
}
),
},
)
is True
)
def test_windows_intel_registry_missing_class_key_is_false(monkeypatch):
monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None))
assert ilp.windows_intel_gpu_in_registry() is False
def _detect_windows_host(
monkeypatch,
winreg_fake,
powershell_stdout = "",
):
"""Drive the real detect_host() as a GPU-less Windows host with a fake
registry, recording every run_capture invocation. Pins the wiring the
unit tests above cannot see: registry-first, CIM only on a registry miss."""
monkeypatch.setitem(sys.modules, "winreg", winreg_fake)
monkeypatch.setattr(ilp.platform, "system", lambda: "Windows")
monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64")
for _env in (
"CUDA_VISIBLE_DEVICES",
"HIP_VISIBLE_DEVICES",
"ROCR_VISIBLE_DEVICES",
"HIP_PATH",
"ROCM_PATH",
):
monkeypatch.delenv(_env, raising = False)
monkeypatch.setattr(
ilp.shutil,
"which",
lambda name: "powershell" if name in ("powershell", "pwsh") else None,
)
captured = []
def _fake_run_capture(command, **kwargs):
captured.append(command[0])
if command[0] == "powershell":
return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "")
return SimpleNamespace(returncode = 1, stdout = "", stderr = "")
monkeypatch.setattr(ilp, "run_capture", _fake_run_capture)
return ilp.detect_host(), captured
def test_detect_host_registry_intel_skips_cim_probe(monkeypatch):
winreg = _FakeWinreg(
_FakeRegKey(
subkeys = {
"0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}),
}
)
)
host, captured = _detect_windows_host(monkeypatch, winreg)
assert host.has_intel_gpu is True
assert "powershell" not in captured
def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch):
winreg = _FakeWinreg(
_FakeRegKey(
subkeys = {
"0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}),
}
)
)
host, captured = _detect_windows_host(
monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics"
)
assert host.has_intel_gpu is True
assert "powershell" in captured
def test_windows_intel_registry_unexpected_error_is_false(monkeypatch):
# The probe is advisory: even a non-OSError bug in the walk must return
# False (deferring to the CIM fallback), never crash detect_host.
class _ExplodingWinreg:
HKEY_LOCAL_MACHINE = object()
def OpenKey(self, parent, name):
raise TypeError(name)
monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg())
assert ilp.windows_intel_gpu_in_registry() is False
def test_detect_host_cim_rescues_exploding_registry(monkeypatch):
class _ExplodingWinreg:
HKEY_LOCAL_MACHINE = object()
def OpenKey(self, parent, name):
raise TypeError(name)
host, captured = _detect_windows_host(
monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics"
)
assert host.has_intel_gpu is True
assert "powershell" in captured

9
studio/backend/tests/test_linux_external_media_paths.py
View file

@ -252,10 +252,17 @@ def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tm
outputs_root = lambda: tmp_path / "missing-outputs",
exports_root = lambda: tmp_path / "missing-exports",
)
fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root])
fake_external_media = SimpleNamespace(
linux_run_media_mount_roots = lambda: [media_root],
windows_drive_roots = lambda: [],
)
fake_studio_db = SimpleNamespace(
list_scan_folders = lambda: [],
contains_sensitive_path_component = studio_db.contains_sensitive_path_component,
# The media root is a legitimate mount, not denied; the .ssh 403 below
# comes from the credential check. A False stub keeps this OS-independent
# (on macOS tmp_path lives under the denied /private/var).
is_denied_system_path = lambda _p: False,
)
monkeypatch.setitem(sys.modules, "utils.paths", fake_paths)
monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media)

320
studio/backend/tests/test_llama_admission.py Normal file
View file

@ -0,0 +1,320 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import asyncio
import os
import sys
import threading
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference import llama_admission
from core.inference.llama_admission import (
ADMISSION_CONTROL_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
DEFAULT_ADMISSION_MAX_QUEUE,
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S,
LlamaAdmissionConfig,
LlamaAdmissionQueueFull,
get_llama_admission_queue,
llama_admission_config_from_env,
reset_llama_admission_queues,
)
@pytest.fixture(autouse = True)
def _reset_queues():
reset_llama_admission_queues()
yield
reset_llama_admission_queues()
def test_admission_config_defaults(monkeypatch):
for name in (
ADMISSION_CONTROL_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
):
monkeypatch.delenv(name, raising = False)
config = llama_admission_config_from_env()
assert config.enabled is True
assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE
def test_admission_config_env_overrides(monkeypatch):
monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off")
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0")
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25")
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0")
config = llama_admission_config_from_env()
assert config.enabled is False
assert config.queue_timeout_s is None
assert config.keepalive_interval_s == 0.25
assert config.max_queue is None
def test_admission_config_positive_queue_timeout_env(monkeypatch):
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600")
config = llama_admission_config_from_env()
assert config.queue_timeout_s == 600.0
def test_fifo_capacity_one_grants_next_waiter_on_release():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
third = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second.lease_nowait() is None
assert third.lease_nowait() is None
assert queue.snapshot().queued == 2
first_lease.release()
second_lease = await second.wait(0.1)
assert second_lease is not None
assert third.lease_nowait() is None
second_lease.release()
third_lease = await third.wait(0.1)
assert third_lease is not None
third_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_queue_full_rejects_excess_waiter():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(max_queue = 1)
first = queue.reserve(capacity = 1, config = config)
queued = queue.reserve(capacity = 1, config = config)
assert first.lease_nowait() is not None
assert queued.lease_nowait() is None
with pytest.raises(LlamaAdmissionQueueFull):
queue.reserve(capacity = 1, config = config)
asyncio.run(_run())
def test_disabled_admission_bypasses_active_slot_limit():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(enabled = False)
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
assert first.lease_nowait() is not None
assert second.lease_nowait() is not None
assert queue.snapshot().active == 0
assert queue.snapshot().queued == 0
asyncio.run(_run())
def test_cancelling_promoted_waiter_releases_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
first_lease.release()
await asyncio.sleep(0)
second.cancel()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_cancelling_promoted_waiter_before_delivery_releases_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
first_lease.release()
second.cancel()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_external_waiter_future_cancel_invalidates_reservation():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second._waiter is not None
second._waiter.future.cancel()
assert second.lease_nowait() is None
assert second.is_cancelled is True
assert await second.wait(0.01) is None
first_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_wait_returns_none_when_waiter_future_cancelled_during_wait():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second._waiter is not None
wait_task = asyncio.create_task(second.wait(1.0))
await asyncio.sleep(0)
second._waiter.future.cancel()
assert await asyncio.wait_for(wait_task, timeout = 0.1) is None
assert second.is_cancelled is True
first_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_capacity_increase_promotes_existing_waiter_fifo():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second.lease_nowait() is None
assert queue.snapshot().active == 1
assert queue.snapshot().queued == 1
third = queue.reserve(capacity = 2, config = config)
second_lease = await second.wait(0.1)
assert second_lease is not None
assert third.lease_nowait() is None
snapshot = queue.snapshot()
assert snapshot.capacity == 2
assert snapshot.active == 2
assert snapshot.queued == 1
first_lease.release()
third_lease = await third.wait(0.1)
assert third_lease is not None
second_lease.release()
third_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_lease_release_is_idempotent_under_concurrent_calls():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
reservation = queue.reserve(capacity = 1, config = config)
lease = reservation.lease_nowait()
assert lease is not None
threads = [threading.Thread(target = lease.release) for _ in range(16)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_new_key_evicts_idle_prior_load_queues():
# Each model load carries a fresh ephemeral port, so a new base_url key must
# not leave the drained queues from earlier loads accumulating forever.
get_llama_admission_queue("http://127.0.0.1:1001")
get_llama_admission_queue("http://127.0.0.1:1002")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"}
get_llama_admission_queue("http://127.0.0.1:1003")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"}
def test_new_key_retains_in_flight_prior_load_queue():
config = LlamaAdmissionConfig()
busy = get_llama_admission_queue("http://127.0.0.1:2001")
async def _run():
reservation = busy.reserve(capacity = 1, config = config)
lease = reservation.lease_nowait()
assert lease is not None
# A new load must not drop a queue that still has an in-flight request.
get_llama_admission_queue("http://127.0.0.1:2002")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"}
# Once it drains, the next load reclaims it.
lease.release()
get_llama_admission_queue("http://127.0.0.1:2003")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"}
asyncio.run(_run())

53
studio/backend/tests/test_llama_cpp_effective_parallel_slots.py Normal file
View file

@ -0,0 +1,53 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import os
import sys
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference import llama_cpp as llama_cpp_module
from core.inference.llama_cpp import LlamaCppBackend
@pytest.fixture
def backend(monkeypatch):
monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0)
monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None)
return LlamaCppBackend()
def test_effective_parallel_slots_initial_value_is_one(backend):
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend):
backend._commit_effective_parallel_slots(3)
assert backend.effective_parallel_slots == 3
@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"])
def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value):
backend._commit_effective_parallel_slots(value)
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_reset_returns_to_one(backend):
backend._commit_effective_parallel_slots(4)
backend._reset_effective_parallel_slots()
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_unload_resets_to_one(backend):
backend._commit_effective_parallel_slots(4)
backend.unload_model()
assert backend.effective_parallel_slots == 1

6
studio/backend/tests/test_llama_cpp_freshness.py
View file

@ -137,9 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path):
@pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"])
def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo):
# The freshness check queries whichever release repo the marker records,
# so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right
# "latest" tag.
# The freshness check queries whichever release repo the marker records:
# new installs record the fork, legacy CPU/macOS markers still say ggml-org,
# and both must get the right "latest" tag.
install_dir = tmp_path / "llama.cpp"
_write_marker(install_dir, tag = "b9000", published_repo = repo)
bin_path = _fake_binary(install_dir, layout = "cmake")

192
studio/backend/tests/test_llama_cpp_stream_cancel.py Normal file
View file

@ -0,0 +1,192 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import contextlib
import os
import socket
import sys
import threading
import time
import httpx
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled
def _backend_stub() -> LlamaCppBackend:
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
backend._port = 48848
backend._effective_context_length = 4096
backend._supports_reasoning = False
backend._reasoning_always_on = False
backend._reasoning_style = "enable_thinking"
backend._supports_preserve_thinking = False
return backend
def test_stream_cancel_uses_internal_exception_not_generator_exit():
class FakeResponse:
status_code = 200
def close(self):
pass
class FakeStream:
def __enter__(self):
return FakeResponse()
def __exit__(self, *_args):
return False
class FakeClient:
def stream(self, *_args, **_kwargs):
return FakeStream()
cancel_event = threading.Event()
with pytest.raises(Exception) as exc_info:
with LlamaCppBackend._stream_with_retry(
FakeClient(),
"http://llama.test/v1/chat/completions",
{},
cancel_event,
):
cancel_event.set()
raise httpx.ReadError("client closed")
assert exc_info.type is _LlamaStreamCancelled
assert not issubclass(exc_info.type, GeneratorExit)
def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch):
backend = _backend_stub()
@contextlib.contextmanager
def fake_open_stream(*_args, **_kwargs):
raise _LlamaStreamCancelled
monkeypatch.setattr(backend, "_open_stream", fake_open_stream)
chunks = list(
backend.generate_chat_completion(
[{"role": "user", "content": "hi"}],
cancel_event = threading.Event(),
)
)
assert chunks == []
class _StallUpstream:
"""Raw HTTP/1.1 server that streams one chunked SSE chunk, then holds the
socket open and silent so the client's next read blocks in recv() until its
side is torn down. Reproduces a mid-stream stall (llama-server goes quiet)."""
def __init__(self):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self._sock.bind(("127.0.0.1", 0))
self._sock.listen(1)
self.port = self._sock.getsockname()[1]
self._stop = threading.Event()
self._thread = threading.Thread(target = self._serve, daemon = True)
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.port}/v1/chat/completions"
def __enter__(self):
self._thread.start()
return self
def __exit__(self, *_exc):
self._stop.set()
try:
self._sock.close()
except OSError:
pass
self._thread.join(timeout = 5)
def _serve(self) -> None:
try:
conn, _ = self._sock.accept()
except OSError:
return
with conn:
conn.settimeout(5)
try:
buf = b""
while b"\r\n\r\n" not in buf:
data = conn.recv(4096)
if not data:
return
buf += data
head, _, body = buf.partition(b"\r\n\r\n")
content_length = 0
for line in head.split(b"\r\n"):
if line.lower().startswith(b"content-length:"):
content_length = int(line.split(b":", 1)[1].strip())
break
while len(body) < content_length:
data = conn.recv(4096)
if not data:
break
body += data
except OSError:
return
conn.sendall(
b"HTTP/1.1 200 OK\r\n"
b"Content-Type: text/event-stream\r\n"
b"Transfer-Encoding: chunked\r\n"
b"\r\n"
)
chunk = b"data: hello\n\n"
conn.sendall(b"%x\r\n%s\r\n" % (len(chunk), chunk))
# Stall: stay open and silent until the client shuts its side down.
while not self._stop.wait(timeout = 0.05):
try:
conn.settimeout(0.05)
if conn.recv(1) == b"":
return
except socket.timeout:
continue
except OSError:
return
def test_cancel_interrupts_a_read_blocked_on_a_mid_stream_stall():
# Mid-stream stall: the reader is parked in recv() on a long bound read timeout,
# so response.close() alone can't wake it; the watcher must shut the socket down.
# Assert cancel lands in seconds, not at the far-off deadline (pre-fix: hung ~30s).
with _StallUpstream() as server:
cancel_event = threading.Event()
def _cancel_soon():
time.sleep(0.3)
cancel_event.set()
threading.Thread(target = _cancel_soon, daemon = True).start()
started = time.monotonic()
with httpx.Client(
limits = httpx.Limits(max_keepalive_connections = 0), trust_env = False
) as client:
with pytest.raises(_LlamaStreamCancelled):
with LlamaCppBackend._stream_with_retry(
client,
server.url,
{},
cancel_event,
first_token_deadline = started + 30,
) as response:
for _chunk in response.iter_text():
pass # first chunk arrives, then the read blocks silently
elapsed = time.monotonic() - started
assert elapsed < 10, f"cancel took {elapsed:.1f}s; the blocked read was not interrupted"

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

@ -1498,6 +1498,59 @@ def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch)
assert any("Let me search." in t for t in content_texts)
def test_textual_explicit_id_reuses_provisional_card(monkeypatch):
# A textual Mistral-style call with an explicit ``id`` must reconcile onto the
# open provisional TEXT card (keyed "call_0"), not spawn a duplicate under the
# explicit id (which the parser keeps for execution).
big_query = "cats " * 80 # push the drained call past the provisional floor
call = "[TOOL_CALLS]" + json.dumps(
[{"name": "web_search", "arguments": {"query": big_query}, "id": "explicit-42"}]
)
assert len(call) > 256
# Small chunks so the provisional card opens mid-generation (a single-shot
# delta parses instantly and never shows a provisional to exercise).
chunks = [call[i : i + 24] for i in range(0, len(call), 24)]
streams = [
[_sse({"content": c}) for c in chunks] + [_done()],
[_sse({"content": "done"}), _done()],
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, streams, payloads)
calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [("web_search", {"query": big_query})]
tool_starts = [e for e in events if e.get("type") == "tool_start"]
# Empty-args card = provisional open; full-args card = reconciled real start.
provisional = [e for e in tool_starts if not e.get("arguments")]
real = [e for e in tool_starts if e.get("arguments", {}).get("query")]
assert len(provisional) == 1, tool_starts # provisional actually opened
prov_id = provisional[0]["tool_call_id"]
# Exactly one real card, sharing the provisional id, not a duplicate under
# the explicit "explicit-42" id.
assert len(real) == 1, tool_starts
assert real[0]["tool_call_id"] == prov_id
assert real[0]["tool_name"] == "web_search"
assert {e["tool_call_id"] for e in tool_starts} == {prov_id}
# A single tool_end reconciles the card; no stale empty-result close.
ends = [e for e in events if e.get("type") == "tool_end"]
assert [e["tool_call_id"] for e in ends] == [prov_id]
assert ends[0]["result"] == "result"
def test_textual_llama_python_tag_marker_not_leaked(monkeypatch):
# Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form.
streams = [
@ -1826,6 +1879,39 @@ def test_large_python_tool_call_emits_early_provisional_start(monkeypatch):
assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events)
def test_auto_mode_render_html_suppresses_provisional_card_under_confirm(monkeypatch):
"""render_html is no longer unconditionally safe (a networked canvas asks), so
with confirm_tool_calls set under permission_mode="auto" its early provisional
card is suppressed; the real full-argument tool_start still fires and a static
canvas runs without a prompt."""
args = {"code": "<html>" + "x" * 80 + "</html>"}
first_stream = _streamed_structured_tool_call("render_html", args, "call_rh")
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda name, arguments, **_k: "OK")
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "make a card"}],
tools = [{"type": "function", "function": {"name": "render_html"}}],
confirm_tool_calls = True,
permission_mode = "auto",
max_tool_iterations = 1,
)
)
tool_starts = [e for e in events if e.get("type") == "tool_start"]
provisional = [e for e in tool_starts if not e.get("arguments")]
# The confirm gate now suppresses the early provisional card for render_html.
assert provisional == [], tool_starts
real = [e for e in tool_starts if e.get("arguments")]
assert real and real[0]["tool_name"] == "render_html"
# A static canvas is classified safe, so it still runs without an approval gate.
assert real[0].get("awaiting_confirmation") in (False, None)
def test_small_python_tool_call_has_no_provisional_start(monkeypatch):
"""A small tool-call argument finishes streaming instantly, so it keeps the
existing behavior of a single (real) tool_start with no provisional card."""
@ -2883,3 +2969,204 @@ def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch):
m.get("role") == "user" and "used all available tool calls" in m.get("content", "")
for m in payloads[2]["messages"]
), payloads[2]["messages"]
# ── Live tool-call argument streaming (tool_args events) ─────────────────────
def _python_tool_schema() -> list[dict]:
return [
{
"type": "function",
"function": {
"name": "python",
"description": "Run python code.",
"parameters": {
"type": "object",
"properties": {"code": {"type": "string"}},
"required": ["code"],
},
},
}
]
def test_structured_tool_args_stream_to_provisional_card(monkeypatch):
"""A large structured tool call must stream its arguments as tool_args events
to the provisional card (backlog that triggered the card, then each
fragment), while the executed call and the model's view stay exactly what the
accumulator built."""
code = "print('x')\n" + ("# pad\n" * 80)
args_json = json.dumps({"code": code})
call_id = "call_live_args"
split = _PROVISIONAL_ARGS_MIN_CHARS + 16
frag1, frag2, frag3 = (
args_json[:split],
args_json[split : split + 40],
args_json[split + 40 :],
)
def _tc_delta(fragment: str, with_header: bool) -> str:
entry: dict = {"index": 0, "function": {"arguments": fragment}}
if with_header:
entry.update({"id": call_id, "type": "function"})
entry["function"]["name"] = "python"
return _sse({"tool_calls": [entry]})
first_stream = [
_tc_delta(frag1, with_header = True),
_tc_delta(frag2, with_header = False),
_tc_delta(frag3, with_header = False),
_done(),
]
second_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
executed: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
executed.append((name, arguments))
return "ok"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run it"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
)
)
starts = [e for e in events if e.get("type") == "tool_start"]
assert starts and starts[0]["tool_call_id"] == call_id
args_events = [e for e in events if e.get("type") == "tool_args"]
assert args_events, "no tool_args events were streamed"
assert all(e["tool_call_id"] == call_id for e in args_events)
# First event is the backlog, the rest raw fragments; together the args JSON.
assert args_events[0]["text"] == frag1
assert "".join(e["text"] for e in args_events) == args_json
# The streamed display path must not perturb execution or the model view.
assert executed == [("python", {"code": code})]
assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"]
tc = assistant_messages[-1]["tool_calls"][0]
assert tc["id"] == call_id
# Controller re-serializes args (normalized JSON); parsed payload unchanged.
assert json.loads(tc["function"]["arguments"]) == {"code": code}
def test_text_tool_call_streams_args_and_reconciles_card(monkeypatch):
"""A TEXT (XML) tool call must stream its raw call text as tool_args under the
id the stream-end parser assigns ("call_0"), so the provisional card and the
final tool_start reconcile."""
code = "print('hello')\n" + ("# filler\n" * 60)
call_json = json.dumps({"name": "python", "arguments": {"code": code}})
call_text = f"<tool_call>{call_json}</tool_call>"
chunks = [call_text[i : i + 48] for i in range(0, len(call_text), 48)]
first_stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
second_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
executed: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
executed.append((name, arguments))
return "ok"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run it"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
)
)
starts = [e for e in events if e.get("type") == "tool_start"]
assert starts, "no tool_start emitted"
# Provisional card first (parser's first-call id), then the reconciling start.
assert starts[0]["tool_call_id"] == "call_0"
assert starts[0]["arguments"] == {}
assert starts[-1]["tool_call_id"] == "call_0"
args_events = [e for e in events if e.get("type") == "tool_args"]
assert args_events, "no tool_args events for the text call"
assert all(e["tool_call_id"] == "call_0" for e in args_events)
streamed = "".join(e["text"] for e in args_events)
# Streamed text is the drained call (display only); it must never leak into
# content events.
assert '"name": "python"' in streamed
assert executed == [("python", {"code": code})]
content_events = [e for e in events if e.get("type") == "content"]
assert not any("<tool_call>" in e["text"] for e in content_events)
def test_ordinary_json_answer_streams_no_tool_args(monkeypatch):
"""A large ordinary JSON answer (no enabled tool name) must not spawn a
provisional card or tool_args events; it stays a normal content answer."""
answer = json.dumps({"result": "fine", "data": ["x" * 40] * 12, "note": "not a tool call"})
chunks = [answer[i : i + 64] for i in range(0, len(answer), 64)]
stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream], payloads)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "give me json"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
)
)
assert not [e for e in events if e.get("type") == "tool_args"]
assert not [e for e in events if e.get("type") == "tool_start"]
content_events = [e for e in events if e.get("type") == "content"]
assert content_events and answer in content_events[-1]["text"]
def test_provisional_text_card_closed_when_parse_fails(monkeypatch):
"""A >=256-char enabled-name text sniff opens a provisional card; if the
drained text then fails to parse (auto-heal off, truncated call), the
DRAINING false-positive path must close the card with a tool_end instead of
leaving it spinning forever."""
# Truncated mid-arguments and never closed: unparseable without healing.
call_text = '<tool_call>{"name": "python", "arguments": {"code": "' + "x" * (
_PROVISIONAL_ARGS_MIN_CHARS + 64
)
chunks = [call_text[i : i + 48] for i in range(0, len(call_text), 48)]
stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream], payloads)
executed: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
executed.append((name, arguments))
return "ok"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run it"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
auto_heal_tool_calls = False,
)
)
starts = [e for e in events if e.get("type") == "tool_start"]
ends = [e for e in events if e.get("type") == "tool_end"]
assert starts and starts[0]["tool_call_id"] == "call_0"
assert executed == [] # nothing parsed, nothing ran
assert ends, "provisional card left dangling (no tool_end)"
assert ends[-1]["tool_call_id"] == "call_0"

130
studio/backend/tests/test_llama_cpp_update.py
View file

@ -393,6 +393,9 @@ def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path):
assert "--llama-tag" in cmd and "latest" in cmd
assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x"
assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd
# No pin: source-build detection and the unpinned apply share the same
# "latest" resolver, so they already agree.
assert "--published-release-tag" not in cmd
def test_start_update_happy_path(monkeypatch, tmp_path):
@ -448,6 +451,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
# A Vulkan install (marker asset carries 'vulkan') must re-assert
# UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to
# CUDA/ROCm and silently replaces the Vulkan build.
install_dir = tmp_path / "llama.cpp"
binary = _write_install(
install_dir,
"b9493",
repo = "ggml-org/llama.cpp",
asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz",
)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
def _on_start(cmd):
_write_install(
install_dir,
"b9518",
repo = "ggml-org/llama.cpp",
asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
)
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
lines = ["installed\n"],
on_start = _on_start,
captured_kwargs = popen_kwargs,
)
assert upd.start_update()["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")
@ -477,6 +522,57 @@ def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
assert "Updated llama.cpp to b9596-mix-e6f2453." in job["message"]
def _run_start_update_to_completion():
res = upd.start_update()
assert res["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
return job
time.sleep(0.05)
return upd.get_update_status()["job"]
def test_start_update_pinned_tag_mismatch_fails(monkeypatch, tmp_path):
# Installer stays on the pinned repo but produces a different tag -> it
# ignored the pin (the silent mismatch this pin exists to prevent). Fail loud.
monkeypatch.setattr(sys, "platform", "linux")
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9601-mix-a0e2906"
)
_patch_installer_popen(
monkeypatch,
on_start = lambda cmd: _write_install(install_dir, "b9500", release_tag = "b9500-mix-deadbee"),
)
job = _run_start_update_to_completion()
assert job["state"] == "error", job
assert "b9601-mix-a0e2906" in (job["error"] or "")
def test_start_update_pinned_reroute_to_other_repo_ok(monkeypatch, tmp_path):
# A Vulkan/Intel host reroutes fork->upstream and drops the pin, installing a
# different-repo tag. Legitimate: the pin check must not flag the repo switch.
monkeypatch.setattr(sys, "platform", "linux")
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595", repo = "unslothai/llama.cpp")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9601-mix-a0e2906"
)
_patch_installer_popen(
monkeypatch,
on_start = lambda cmd: _write_install(install_dir, "b9601", repo = "ggml-org/llama.cpp"),
)
job = _run_start_update_to_completion()
assert job["state"] == "success", job
def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
@ -594,9 +690,10 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path):
def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path):
# CPU installs come from ggml-org. Re-running into the same install-dir/repo
# reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU
# detection) is reserved for setup.sh's arm64 rescue and must not appear here.
# Legacy CPU installs recorded a ggml-org marker (new installs use the fork).
# Re-running into the same install-dir/repo reproduces the same CPU bundle;
# --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's
# arm64 rescue and must not appear here.
cmd = _capture_install_cmd(
monkeypatch,
tmp_path,
@ -620,6 +717,33 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm
assert "--cpu-fallback" not in cmd
def test_install_cmd_pins_offered_release_tag(monkeypatch, tmp_path):
# Apply must install exactly the release the banner offered. The installer's
# own "latest" comes from commit-date-ordered sources, which can lag the
# published_at-newest tag detection picked; unpinned, that lag makes Update
# reinstall the current build while the banner never clears.
monkeypatch.setattr(sys, "platform", "linux")
cmd = _capture_install_cmd(monkeypatch, tmp_path, latest = "b9601-mix-a0e2906")
# The full release identity is pinned, not the bare upstream base.
assert cmd[cmd.index("--published-release-tag") + 1] == "b9601-mix-a0e2906"
def test_install_cmd_pins_on_windows(monkeypatch, tmp_path):
# The darwin exemption must not leak to other platforms.
monkeypatch.setattr(sys, "platform", "win32")
cmd = _capture_install_cmd(monkeypatch, tmp_path)
assert cmd[cmd.index("--published-release-tag") + 1] == "b9518"
def test_install_cmd_does_not_pin_on_macos(monkeypatch, tmp_path):
# A pinned tag disables the installer's older-release walk-back, which macOS
# needs to skip prebuilts built for a newer macOS than the host.
monkeypatch.setattr(sys, "platform", "darwin")
cmd = _capture_install_cmd(monkeypatch, tmp_path)
assert "--published-release-tag" not in cmd
assert "--llama-tag" in cmd and "latest" in cmd
# --- refusal + maintenance-state coordination ---

193
studio/backend/tests/test_llama_cpp_vulkan_probe.py Normal file
View file

@ -0,0 +1,193 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Vulkan free-VRAM reader regression tests on a synthetic probe output.
Covers the post-probe handling in
``LlamaCppBackend._get_gpu_free_memory_vulkan``:
* integrated GPUs (probe reports is_igpu=1) leave a flat per-device host
margin matching llama.cpp's --fit-target, so context auto-sizing can't
over-commit shared RAM, and report total 0 (shared RAM is not a budget),
* discrete GPUs (is_igpu=0) keep their free untouched and pass their real
total through so the fit can reserve absolute headroom,
* an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged
(ggml applies it), not stripped or filtered in Python -- the probe reports
ggml's compact ordinal, which load_model pins with ``--device Vulkan<i>``.
The ggml Vulkan library is never loaded: subprocess.run is mocked to emit
the tab-separated lines the real ``_vulkan_probe.py`` would print.
"""
from __future__ import annotations
import subprocess
import sys
import types as _types
from pathlib import Path
from unittest import mock
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import importlib as _importlib # noqa: E402
def _maybe_stub(name: str, builder):
try:
_importlib.import_module(name)
except ImportError:
sys.modules[name] = builder()
def _build_loggers_stub():
m = _types.ModuleType("loggers")
m.get_logger = lambda name: __import__("logging").getLogger(name)
return m
_maybe_stub("loggers", _build_loggers_stub)
_maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
from core.inference import llama_cpp as _llama_mod # noqa: E402
from core.inference.llama_cpp import ( # noqa: E402
LlamaCppBackend,
_llama_lib_dir,
_vulkan_lib_filename,
)
MIB = 1024 * 1024
GIB = 1024 * MIB
def _make_vulkan_install(tmp_path: Path) -> str:
"""A binary whose sibling dir holds the Vulkan ggml lib, so the
reader's ``is_vulkan_backend`` sibling-file check passes."""
bindir = tmp_path / "build" / "bin"
bindir.mkdir(parents = True)
binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server")
binary.write_bytes(b"stub")
(bindir / _vulkan_lib_filename()).write_bytes(b"stub")
return str(binary)
def _mock_probe(rows: list[str], captured_env: dict | None = None):
"""Patch subprocess.run so the _vulkan_probe.py call returns ``rows``
(already tab-formatted), recording the env it was launched with."""
real_run = subprocess.run
def fake_run(cmd, *args, **kwargs):
if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd):
if captured_env is not None:
captured_env.clear()
captured_env.update(kwargs.get("env") or {})
return subprocess.CompletedProcess(
args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = ""
)
return real_run(cmd, *args, **kwargs)
return mock.patch("subprocess.run", side_effect = fake_run)
def _row(
idx: int,
free_bytes: int,
is_igpu: int,
total_bytes: int = 0,
) -> str:
return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
def test_integrated_gpu_leaves_host_margin(tmp_path):
binary = _make_vulkan_install(tmp_path)
# iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target).
# total stays 0: shared system RAM is not a VRAM budget for the fit.
rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus
def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path):
binary = _make_vulkan_install(tmp_path)
# 6 GiB free on a partially occupied 24 GiB card: free is untouched and the
# real total flows through so the fit reserves absolute headroom (CUDA/ROCm
# parity) instead of the looser free*frac budget.
rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus
def test_large_discrete_gpu_is_untouched(tmp_path):
binary = _make_vulkan_install(tmp_path)
# A 48 GiB discrete card stays untouched regardless of size; only the
# iGPU flag triggers the host margin, never a VRAM/RAM ratio.
rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus
def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch):
# The mask is NOT stripped or filtered in Python: ggml parses it in raw
# physical-device space while this probe reports the compact post-filter
# ordinal, so mixing spaces would be wrong. It is passed through unchanged
# so ggml applies it to the same device list the launch will enumerate.
binary = _make_vulkan_install(tmp_path)
monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1")
captured: dict = {}
rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows, captured_env = captured):
LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured
def test_vulkan_pin_args_uses_device_names_not_env_mask():
# Pin by compact device name via --device (the space the probe reports and
# the registry names), never by writing a compact ordinal into the raw
# GGML_VK_VISIBLE_DEVICES index space.
assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"]
assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"]
assert LlamaCppBackend._vulkan_pin_args(None) == []
assert LlamaCppBackend._vulkan_pin_args([]) == []
def test_vulkan_only_build_is_detected(tmp_path):
binary = _make_vulkan_install(tmp_path)
assert LlamaCppBackend._is_vulkan_backend(binary) is True
def test_multi_backend_build_is_not_vulkan_only(tmp_path):
# A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be
# treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan
# device; defer to the CUDA/HIP path instead.
binary = _make_vulkan_install(tmp_path)
cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so"
(_llama_lib_dir(binary) / cuda).write_bytes(b"stub")
assert LlamaCppBackend._is_vulkan_backend(binary) is False
@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX")
def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path):
# create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root
# when it cannot symlink; _find_llama_server_binary returns that root entrypoint,
# so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else
# _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently
# never engage on a valid Vulkan install.
import os
binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib
bindir = Path(binary).parent
wrapper = tmp_path / "llama-server"
wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n')
os.chmod(wrapper, 0o755)
assert _llama_lib_dir(str(wrapper)) == bindir
assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))

84
studio/backend/tests/test_llama_route_timeouts.py
View file

@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel():
asyncio.run(_run(False))
asyncio.run(_run(True))
def test_stream_stall_timeout_callable_re_resolved_each_read():
# The OpenAI passthrough passes a callable so the stall bound can switch to
# the short post-terminal grace mid-stream; it must be re-resolved per read,
# not captured once at generator start.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
values = iter([100.0, 2.0])
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 3:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 1,
post_first_item_read_timeout_s = lambda: next(values, 5.0),
):
seen.append(response.request.extensions["timeout"].get("read"))
assert len(seen) == 3
# The callable is resolved right after the first item (arming the
# post-first window) and again before each later read, consuming
# successive values.
assert seen[0] == 100.0
assert 1.0 <= seen[1] <= 2.0
assert 4.0 <= seen[2] <= 5.0
asyncio.run(_run())
def test_stream_stall_timeout_disabled_clears_read_timeout():
# UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so
# the callable returns None. Once a chunk has arrived the leftover
# first-token read timeout must be cleared, else a long post-first-chunk gap
# trips a stale deadline the operator asked to turn off.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 2:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 5,
post_first_item_read_timeout_s = lambda: None,
):
seen.append(response.request.extensions["timeout"].get("read"))
# The first-token path armed a finite read timeout; after the first chunk
# with the guard disabled, it is cleared to None on every subsequent read.
assert seen == [None, None], seen
asyncio.run(_run())

48
studio/backend/tests/test_load_progress_throttle.py Normal file
View file

@ -0,0 +1,48 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The /api/inference/load-progress throttle: one line per 10% step, reset per load."""
import pytest
import routes.inference as ri
class _Capture:
def __init__(self):
self.events = []
def info(self, event, **kw):
self.events.append((event, kw))
@pytest.fixture
def cap(monkeypatch):
capture = _Capture()
monkeypatch.setattr(ri, "logger", capture)
ri._reset_load_progress_step()
return capture
def _percents(cap):
return [kw["percent"] for _event, kw in cap.events]
def test_new_load_first_step_logs_after_reset(cap):
# Load A reaches 100%.
ri._log_load_progress_step(1.0, "ready")
assert _percents(cap) == [100]
# Same value keeps deduping (steady poll on a finished load stays quiet).
ri._log_load_progress_step(1.0, "ready")
assert _percents(cap) == [100]
# A new load arms the throttle, so a cached load B that reports 100% on its
# first poll still emits its progress line instead of hitting step == prev.
ri._reset_load_progress_step()
ri._log_load_progress_step(1.0, "ready")
assert _percents(cap) == [100, 100]
def test_steady_poll_dedups_within_a_load(cap):
for _ in range(3):
ri._log_load_progress_step(0.3, "mmap")
assert _percents(cap) == [30] # one line per 10% step, not one per poll

121
studio/backend/tests/test_logging_middleware.py
View file

@ -135,7 +135,7 @@ def test_duplicate_get_within_window_deduped(logs, monkeypatch):
mw = LoggingMiddleware(app)
for _ in range(3):
_run(mw(_http_scope("/api/chat/projects"), _noop_receive, send))
_run(mw(_http_scope("/api/models/browse-folders"), _noop_receive, send))
# Only the first of the identical GET/200 burst is logged.
assert len(logs.events) == 1
@ -183,11 +183,11 @@ def test_quiet_poll_paths_use_longer_heartbeat_window(logs, monkeypatch):
for _ in range(3):
_run(mw(_http_scope("/api/inference/monitor"), _noop_receive, send)) # quiet
for _ in range(3):
_run(mw(_http_scope("/api/chat/projects"), _noop_receive, send)) # normal
_run(mw(_http_scope("/api/models/browse-folders"), _noop_receive, send)) # normal
paths = [e[2]["path"] for e in logs.events]
assert paths.count("/api/inference/monitor") == 1 # collapsed to one heartbeat
assert paths.count("/api/chat/projects") == 3 # base dedup off -> all logged
assert paths.count("/api/models/browse-folders") == 3 # base dedup off -> all logged
def test_distinct_query_strings_are_not_deduped(logs, monkeypatch):
@ -242,3 +242,118 @@ def test_fastapi_static_asset_success_skips_log(tmp_path, logs):
assert response.status_code == 200
assert response.text == "body { color: black; }"
assert len(logs.events) == log_count
def _status_app(status):
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": status, "headers": []})
await send({"type": "http.response.body", "body": b""})
return app
async def _drop(message):
pass
def _paths_logged(logs):
return [e[2]["path"] for e in logs.events]
def test_quiet_success_get_2xx_suppressed(logs):
# A GET/2xx poll on a quiet-success path logs nothing; the signal is in events.
for path in ("/api/chat/threads", "/api/export/status", "/api/hub/download-status"):
_run(LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop))
assert logs.events == []
def test_chat_detail_and_message_reads_still_log(logs):
# Only the exact list polls are suppressed; detail/message reads carry latency
# signal and keep their access line.
for path in (
"/api/chat/threads/abc123",
"/api/chat/threads/abc123/messages",
"/api/chat/threads/abc123/messages/m1",
"/api/chat/projects/p1",
):
_run(LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop))
assert _paths_logged(logs) == [
"/api/chat/threads/abc123",
"/api/chat/threads/abc123/messages",
"/api/chat/threads/abc123/messages/m1",
"/api/chat/projects/p1",
]
def test_quiet_success_is_get_only(logs):
# Mutations on the same paths still log (suppression is GET-only).
for method in ("POST", "PUT", "DELETE"):
_run(
LoggingMiddleware(_status_app(200))(
_http_scope("/api/chat/threads", method = method), _noop_receive, _drop
)
)
assert len(logs.events) == 3
def test_chat_pre_auth_401_suppressed_other_errors_logged(logs):
# The transient bootstrap 401 on a chat list GET is dropped, but a 500 (or any
# other status) still logs so real failures stay visible.
_run(
LoggingMiddleware(_status_app(401))(_http_scope("/api/chat/projects"), _noop_receive, _drop)
)
assert logs.events == []
_run(
LoggingMiddleware(_status_app(500))(_http_scope("/api/chat/projects"), _noop_receive, _drop)
)
assert _paths_logged(logs) == ["/api/chat/projects"]
def test_chat_401_logged_after_first_auth_refresh(logs):
# A chat 401 before any successful token refresh is the bootstrap race and is
# dropped, but once /api/auth/refresh has succeeded on this instance later chat
# 401s are real failures and stay visible.
responses: dict[tuple[str, str], int] = {}
async def app(scope, receive, send):
status = responses.get((scope["method"], scope["path"]), 200)
await send({"type": "http.response.start", "status": status, "headers": []})
await send({"type": "http.response.body", "body": b""})
mw = LoggingMiddleware(app)
responses[("GET", "/api/chat/threads")] = 401
_run(mw(_http_scope("/api/chat/threads"), _noop_receive, _drop))
assert logs.events == [] # bootstrap race: suppressed
# A successful refresh (POST, always logged) closes the bootstrap window.
responses[("POST", "/api/auth/refresh")] = 200
_run(mw(_http_scope("/api/auth/refresh", method = "POST"), _noop_receive, _drop))
assert _paths_logged(logs) == ["/api/auth/refresh"]
# Now the same chat 401 is a real failure and logs.
_run(mw(_http_scope("/api/chat/threads"), _noop_receive, _drop))
assert _paths_logged(logs) == ["/api/auth/refresh", "/api/chat/threads"]
def test_export_status_error_still_logs(logs):
# 2xx suppressed, but an HTTP-level error on export status remains visible.
_run(
LoggingMiddleware(_status_app(200))(_http_scope("/api/export/status"), _noop_receive, _drop)
)
assert logs.events == []
_run(
LoggingMiddleware(_status_app(500))(_http_scope("/api/export/status"), _noop_receive, _drop)
)
assert _paths_logged(logs) == ["/api/export/status"]
def test_legacy_download_progress_heartbeats_not_suppressed(logs, monkeypatch):
# Legacy /api/models download polls emit no progress events, so they heartbeat
# (first hit logs, the burst collapses) rather than vanish entirely.
monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 0)
monkeypatch.setattr(hmod, "_QUIET_POLL_DEDUP_MS", 1000)
mw = LoggingMiddleware(_status_app(200))
for _ in range(3):
_run(mw(_http_scope("/api/models/download-progress"), _noop_receive, _drop))
assert _paths_logged(logs) == ["/api/models/download-progress"]

177
studio/backend/tests/test_mcp_flatten_result.py Normal file
View file

@ -0,0 +1,177 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import contextlib
import json
import sys
from pathlib import Path
from types import SimpleNamespace
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference import mcp_client
from core.inference.mcp_client import (
MAX_IMAGE_PAYLOAD_CHARS,
MCP_IMAGES_SENTINEL,
_flatten_result,
call_tool_sync,
)
from core.inference.tool_loop_controller import is_tool_error, strip_result_for_model
PNG_B64 = "iVBORw0KGgoAAAANSUhEUg=="
def _text(value: str) -> SimpleNamespace:
return SimpleNamespace(type = "text", text = value)
def _image(data: str = PNG_B64, mime: str = "image/png") -> SimpleNamespace:
return SimpleNamespace(type = "image", data = data, mimeType = mime)
def _result(
*blocks,
is_error = False,
structured = None,
) -> SimpleNamespace:
return SimpleNamespace(
content = list(blocks),
is_error = is_error,
structured_content = structured,
)
def test_text_only_result_unchanged():
assert _flatten_result(_result(_text("hello"))) == "hello"
def test_image_only_result_keeps_image_and_notes_model():
flat = _flatten_result(_result(_image()))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert body == "[1 image attached; displayed to the user]"
assert json.loads(payload) == [{"data": PNG_B64, "mimeType": "image/png"}]
def test_text_plus_image_keeps_both():
flat = _flatten_result(_result(_text("Took a screenshot"), _image()))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert body == "Took a screenshot\n[1 image attached; displayed to the user]"
assert json.loads(payload)[0]["mimeType"] == "image/png"
def test_multiple_images_pluralized():
flat = _flatten_result(_result(_image(), _image(mime = "image/jpeg")))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert "[2 images attached; displayed to the user]" in body
assert [img["mimeType"] for img in json.loads(payload)] == ["image/png", "image/jpeg"]
def test_strip_result_for_model_drops_image_payload():
flat = _flatten_result(_result(_text("Took a screenshot"), _image()))
stripped = strip_result_for_model(flat)
assert stripped == "Took a screenshot\n[1 image attached; displayed to the user]"
assert PNG_B64 not in stripped
def test_strip_preserves_literal_mcp_sentinel_in_text():
# A tool that legitimately returns text containing the marker (e.g. reading
# source/docs that quote it) must not be truncated: the suffix is not a
# valid JSON image array.
text = "before\n__MCP_IMAGES__: literal from source\nafter"
assert strip_result_for_model(text) == text
def test_strip_preserves_non_image_json_after_marker():
text = 'log line\n__MCP_IMAGES__:["not", "image", "dicts"]'
assert strip_result_for_model(text) == text
def test_strip_removes_only_valid_terminal_envelope():
text = (
"Earlier mention: __MCP_IMAGES__: is documented here"
"\n[1 image attached; displayed to the user]"
'\n__MCP_IMAGES__:[{"data": "AAAA", "mimeType": "image/png"}]'
)
assert strip_result_for_model(text) == (
"Earlier mention: __MCP_IMAGES__: is documented here"
"\n[1 image attached; displayed to the user]"
)
def test_strip_still_handles_images_and_rag_sentinels():
assert strip_result_for_model("output\n__IMAGES__:['a.png']") == "output"
assert strip_result_for_model("answer\n__RAG_SOURCES__:[{}]") == "answer"
def test_error_result_keeps_error_prefix_and_images():
flat = _flatten_result(_result(_text("boom"), _image(), is_error = True))
assert flat.startswith("Error: boom")
assert is_tool_error(flat)
assert MCP_IMAGES_SENTINEL in flat
def test_image_only_error_no_longer_reports_no_content():
flat = _flatten_result(_result(_image(), is_error = True))
assert flat.startswith("Error: [1 image attached")
assert "tool returned no content" not in flat
def test_oversized_image_omitted_with_note():
huge = "A" * (MAX_IMAGE_PAYLOAD_CHARS + 1)
flat = _flatten_result(_result(_image(data = huge)))
assert flat == "[1 image omitted (too large)]"
assert MCP_IMAGES_SENTINEL not in flat
def test_oversized_budget_shared_across_images():
big = "A" * (MAX_IMAGE_PAYLOAD_CHARS - 10)
flat = _flatten_result(_result(_image(data = big), _image()))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert "1 image attached" in body
assert "1 image omitted (too large)" in body
images = json.loads(payload)
assert len(images) == 1 and images[0]["data"] == big
def test_non_image_binary_block_still_ignored():
flat = _flatten_result(
_result(SimpleNamespace(type = "audio", data = PNG_B64, mimeType = "audio/wav"))
)
assert flat == ""
def test_structured_content_fallback_still_used():
flat = _flatten_result(_result(structured = {"ok": True}))
assert flat == "{'ok': True}"
def test_call_tool_sync_passes_raise_on_error_false_and_keeps_error_images(monkeypatch):
# Guards that call_tool_sync passes raise_on_error=False, so an is_error result
# with image content reaches _flatten_result instead of FastMCP raising ToolError.
seen = {}
class _FakeClient:
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
seen["raise_on_error"] = raise_on_error
return _result(_text("boom"), _image(), is_error = True)
@contextlib.asynccontextmanager
async def _fake_client(url, headers, use_oauth):
yield _FakeClient()
monkeypatch.setattr(mcp_client, "_client", _fake_client)
out = call_tool_sync("http://x", None, "take_screenshot", {})
assert seen["raise_on_error"] is False
assert out.startswith("Error: boom")
assert MCP_IMAGES_SENTINEL in out
assert is_tool_error(out)

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

@ -198,7 +198,12 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
async def __aexit__(self, *args):
return False
async def call_tool(self, name, args):
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
import asyncio as _asyncio
await _asyncio.sleep(30) # never finishes during the test
@ -520,7 +525,12 @@ def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch):
async def __aexit__(self, *args):
return False
async def call_tool(self, name, args):
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
return "ran"
monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
@ -787,6 +797,68 @@ def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch):
assert mcp_client.get_cached_tools("s1") == cached
def test_update_rename_keeps_stdio_session(tmp_path, monkeypatch):
"""The edit dialog resends url/headers/oauth unchanged on a rename, so gating
the close on field presence would drop the live stdio session. Only a real
endpoint/auth change may close it."""
import asyncio
import json
_reset_db(tmp_path, monkeypatch)
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
closed: list = []
monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True)
monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a))
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx demo-server",
headers_json = json.dumps({"API_KEY": "x"}),
is_enabled = True,
)
asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(
display_name = "B",
url = "npx demo-server",
headers = {"API_KEY": "x"},
use_oauth = False,
),
current_subject = "u",
)
)
assert closed == []
assert mcp_servers_db.get_server("s1")["display_name"] == "B"
def test_update_stdio_command_change_closes_session(tmp_path, monkeypatch):
"""A real command change must still close the old stdio session."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
closed: list = []
monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True)
monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a))
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx demo-server",
is_enabled = True,
)
asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "npx other-server"), current_subject = "u"
)
)
assert len(closed) == 1
def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch):
"""Disabling a server must drop its cached tools, not leave them unread."""
import asyncio

7
studio/backend/tests/test_mcp_stdio_pr5863.py
View file

@ -93,7 +93,12 @@ class _RecordingClient:
async def list_tools(self):
return [_FakeTool("list_directory"), _FakeTool("write_file")]
async def call_tool(self, name, args):
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
return _FakeResult(f"called {name}")

629
studio/backend/tests/test_mcp_stdio_sessions.py Normal file
View file

@ -0,0 +1,629 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import asyncio
import sys
import threading
import time
from pathlib import Path
from types import SimpleNamespace
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference import mcp_client
from core.inference.mcp_client import call_tool_sync, close_stdio_sessions
STDIO_URL = "npx fake-stateful-server"
HTTP_URL = "https://mcp.example.test/mcp"
def _result(text: str) -> SimpleNamespace:
return SimpleNamespace(
content = [SimpleNamespace(type = "text", text = text)],
is_error = False,
structured_content = None,
)
class FakeClient:
instances: list["FakeClient"] = []
def __init__(self, url: str):
self.url = url
self.entered = 0
self.exited = 0
self.calls: list[tuple[str, dict]] = []
self.connected = False
self.fail_next = False
self.call_delay = 0.0
# Models a dead stdio transport: real Client.is_connected() stays True
# after the subprocess dies, so liveness is probed via the transport.
self.dead = False
self.transport = SimpleNamespace(_is_session_dead = lambda: self.dead)
FakeClient.instances.append(self)
async def __aenter__(self):
self.entered += 1
self.connected = True
return self
async def __aexit__(self, *exc):
self.exited += 1
self.connected = False
def is_connected(self) -> bool:
return self.connected
async def call_tool(self, name: str, args: dict):
if self.call_delay:
await asyncio.sleep(self.call_delay)
if self.fail_next:
self.fail_next = False
self.connected = False
raise RuntimeError("transport closed")
self.calls.append((name, args))
return _result(f"call-{len(self.calls)}")
@pytest.fixture
def fake_clients(monkeypatch):
FakeClient.instances = []
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: FakeClient(url)
)
yield FakeClient.instances
close_stdio_sessions()
def test_stdio_call_without_scope_is_one_shot(fake_clients):
r1 = call_tool_sync(STDIO_URL, None, "browser_navigate", {"url": "https://x.test"})
r2 = call_tool_sync(STDIO_URL, None, "browser_take_screenshot", {})
assert r1 == "call-1"
assert r2 == "call-1"
assert len(fake_clients) == 2
assert all(client.entered == 1 and client.exited == 1 for client in fake_clients)
def test_stdio_sessions_keyed_by_url_and_env(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {})
call_tool_sync("npx other-server", None, "t", {})
call_tool_sync(STDIO_URL, {"ENV_VAR": "1"}, "t", {})
assert len(fake_clients) == 3
def test_stdio_sessions_scoped_per_chat(fake_clients):
# Two conversations must not share one stateful server process.
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat-a")
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat-b")
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat-a")
assert len(fake_clients) == 2
assert fake_clients[0].calls and len(fake_clients[0].calls) == 2
def test_dead_stdio_session_recovers(fake_clients):
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
# Subprocess dies between calls: the dead transport is detected before the
# next dispatch, so the call reconnects on a fresh session instead of failing.
fake_clients[0].dead = True
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 2
assert fake_clients[0].exited == 1
def test_tool_error_does_not_recycle_session(fake_clients, monkeypatch):
from fastmcp.exceptions import ToolError
class ToolFailure(FakeClient):
async def call_tool(self, name, args):
if name == "boom":
raise ToolError("tool exploded") # tool-level: session stays connected
return await super().call_tool(name, args)
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: ToolFailure(url)
)
assert call_tool_sync(STDIO_URL, None, "boom", {}, scope = "chat").startswith("Error: MCP tool")
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 1
def test_http_stays_one_shot(fake_clients):
call_tool_sync(HTTP_URL, None, "t", {})
call_tool_sync(HTTP_URL, None, "t", {})
assert len(fake_clients) == 2
assert all(c.entered == 1 and c.exited == 1 for c in fake_clients)
def test_timeout_discards_stdio_session(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
key = mcp_client._session_key(STDIO_URL, None, "chat")
fake_clients[0].call_delay = 0.5
out = call_tool_sync(
STDIO_URL,
None,
"slow",
{},
timeout = 0.05,
cancel_event = threading.Event(),
scope = "chat",
)
assert "timed out" in out
assert fake_clients[0].exited == 1
assert key not in mcp_client._stdio_key_locks
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 2
def test_no_timeout_allows_long_call(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
fake_clients[0].call_delay = 0.2
# timeout=None means no deadline: the call must not be treated as wedged.
assert call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat") == "call-2"
def test_connect_races_cancel_event(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(5.0)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
ev = threading.Event()
threading.Timer(0.1, ev.set).start()
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, cancel_event = ev)
assert out == "Error: MCP tool 't' cancelled"
assert time.monotonic() - start < 3.0
assert mcp_client._stdio_sessions == {}
def test_connect_respects_caller_timeout(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(5.0)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.2)
assert "timed out" in out
assert time.monotonic() - start < 3.0
assert mcp_client._stdio_sessions == {}
def test_connect_failure_timeout_surfaces_immediately(fake_clients, monkeypatch):
class InitTimeout(FakeClient):
async def __aenter__(self):
raise asyncio.TimeoutError # e.g. fastmcp's own init timeout
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: InitTimeout(url)
)
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 30.0)
assert "timed out" in out
# Must fail fast, not wait out the 30s/60s connect window.
assert time.monotonic() - start < 5.0
assert mcp_client._stdio_sessions == {}
def test_key_lock_wait_honors_cancel_and_timeout(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(1.5)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
first = threading.Thread(target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat"))
first.start()
key = mcp_client._session_key(STDIO_URL, None, "chat")
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
key_lock = mcp_client._stdio_key_locks.get(key)
if key_lock is not None and key_lock.lock.locked():
break
time.sleep(0.01)
# Second same-scope call is stuck behind the first slow connect: Stop must
# interrupt the key-lock wait, and a short tool timeout must bound it.
ev = threading.Event()
threading.Timer(0.2, ev.set).start()
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, cancel_event = ev, scope = "chat")
assert out == "Error: MCP tool 't' cancelled"
assert time.monotonic() - start < 1.0
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.2, scope = "chat")
assert "timed out" in out
assert time.monotonic() - start < 1.0
first.join(10.0)
assert not first.is_alive()
def test_cancel_pre_set_spawns_nothing(fake_clients):
ev = threading.Event()
ev.set()
out = call_tool_sync(STDIO_URL, None, "t", {}, cancel_event = ev)
assert out == "Error: MCP tool 't' cancelled"
assert fake_clients == []
def test_idle_reap_closes_session(fake_clients, monkeypatch):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
key = mcp_client._session_key(STDIO_URL, None, "chat")
assert key in mcp_client._stdio_key_locks
monkeypatch.setattr(mcp_client, "_STDIO_SESSION_IDLE_TTL", 0.0)
mcp_client._reap_idle_stdio_sessions()
assert fake_clients[0].exited == 1
assert mcp_client._stdio_sessions == {}
assert key not in mcp_client._stdio_key_locks
# Next call transparently opens a fresh session.
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 2
def test_reap_skips_in_flight_session(fake_clients, monkeypatch):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
monkeypatch.setattr(mcp_client, "_STDIO_SESSION_IDLE_TTL", 0.0)
session = next(iter(mcp_client._stdio_sessions.values()))
with mcp_client._stdio_sessions_lock:
session.in_flight = 1
try:
mcp_client._reap_idle_stdio_sessions()
assert fake_clients[0].exited == 0
finally:
with mcp_client._stdio_sessions_lock:
session.in_flight = 0
def test_close_during_connect_is_not_cached(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(0.5)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
results: list[str] = []
worker = threading.Thread(
target = lambda: results.append(call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat"))
)
worker.start()
deadline = time.monotonic() + 5.0
while not fake_clients and time.monotonic() < deadline:
time.sleep(0.01)
assert fake_clients # connect is in progress
# Server deleted/updated mid-connect: the session must not be cached after.
close_stdio_sessions(STDIO_URL)
worker.join(10.0)
assert results and results[0].startswith("Error: MCP tool 't' failed")
assert mcp_client._stdio_sessions == {}
assert fake_clients[0].exited == 1
def test_connect_abort_race_still_closes_client(fake_clients, monkeypatch):
class WinsRace(FakeClient):
async def __aenter__(self):
try:
await asyncio.sleep(5.0)
except asyncio.CancelledError:
pass # connect finishes just as the abort lands
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: WinsRace(url))
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.1)
assert "timed out" in out
assert fake_clients[0].entered == 1
assert fake_clients[0].exited == 1 # no orphaned subprocess
assert mcp_client._stdio_sessions == {}
def test_close_unblocks_no_limit_call(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
fake_clients[0].call_delay = 30.0
results: list[str] = []
worker = threading.Thread(
target = lambda: results.append(
call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat")
)
)
worker.start()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
with mcp_client._stdio_sessions_lock:
if session.in_flight >= 1:
break
time.sleep(0.01)
# Server deleted while a no-limit call is in flight: the request thread
# must not hang forever on the stopped session loop.
close_stdio_sessions(STDIO_URL)
worker.join(5.0)
assert not worker.is_alive()
assert results and results[0].startswith("Error: MCP tool 'slow' failed")
def test_lock_wait_timeout_spares_the_borrowed_session(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
fake_clients[0].call_delay = 1.0
results: list[str] = []
slow = threading.Thread(
target = lambda: results.append(
call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat")
)
)
slow.start()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
with mcp_client._stdio_sessions_lock:
if session.in_flight >= 1:
break
time.sleep(0.01)
# A second same-scope call times out waiting for the call lock; it never
# touched the transport, so the shared session must stay alive and cached.
out = call_tool_sync(STDIO_URL, None, "fast", {}, timeout = 0.05, scope = "chat")
assert "timed out" in out
assert fake_clients[0].exited == 0
slow.join(10.0)
assert results == ["call-2"]
assert fake_clients[0].exited == 0
assert len(mcp_client._stdio_sessions) == 1
def test_stale_session_close_deferred_until_borrower_drains(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
fake_clients[0].call_delay = 0.8
results: list[str] = []
slow = threading.Thread(
target = lambda: results.append(
call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat")
)
)
slow.start()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
with mcp_client._stdio_sessions_lock:
if session.in_flight >= 1:
break
time.sleep(0.01)
# The subprocess "dies" mid-call: a new caller replaces the stale session,
# but its close must wait for the slow borrower instead of killing its call.
fake_clients[0].connected = False
out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
assert out == "call-1"
assert len(fake_clients) == 2
assert fake_clients[0].exited == 0
slow.join(10.0)
assert results == ["call-2"]
assert fake_clients[0].exited == 1 # last borrower performed the deferred close
assert len(mcp_client._stdio_sessions) == 1
def test_error_on_closed_session_does_not_retry(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
# A close can surface at the borrower as a plain transport error instead
# of _SessionClosed; that must not be treated as a crash and retried.
fake_clients[0].fail_next = True
session.closed.set()
out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
assert out == "Error: MCP tool 't' failed: MCP server was updated or removed during the call"
assert len(fake_clients) == 1 # no respawn for the removed config
def test_config_check_blocks_stale_publish(fake_clients):
# Simulates a caller that read the server row before an update/delete:
# the row re-check runs after connect and must block caching.
out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat", config_check = lambda: False)
assert out.startswith("Error: MCP tool 't' failed")
assert mcp_client._stdio_sessions == {}
assert fake_clients[0].exited == 1
def test_close_generation_keys_hold_no_secrets(fake_clients):
secret_url = "npx server --token sk-url-secret"
close_stdio_sessions(secret_url, {"API_KEY": "sk-env-secret"})
close_stdio_sessions(secret_url)
gen_keys = list(mcp_client._stdio_cfg_close_gen) + list(mcp_client._stdio_url_close_gen)
assert gen_keys
# These maps are never pruned: neither command/URL nor env may persist.
assert all("sk-url-secret" not in repr(k) and "sk-env-secret" not in repr(k) for k in gen_keys)
def test_overlapping_calls_serialize_on_shared_session(fake_clients, monkeypatch):
class OverlapDetect(FakeClient):
active = 0
max_active = 0
async def call_tool(self, name, args):
OverlapDetect.active += 1
OverlapDetect.max_active = max(OverlapDetect.max_active, OverlapDetect.active)
try:
await asyncio.sleep(0.2)
return await super().call_tool(name, args)
finally:
OverlapDetect.active -= 1
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: OverlapDetect(url)
)
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
workers = [
threading.Thread(target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat"))
for _ in range(2)
]
for worker in workers:
worker.start()
for worker in workers:
worker.join(10.0)
# A stateful server must never see interleaved same-scope operations.
assert OverlapDetect.max_active == 1
assert len(fake_clients) == 1
def test_timeout_budget_spans_connect_and_call(fake_clients, monkeypatch):
class SlowBoth(FakeClient):
async def __aenter__(self):
await asyncio.sleep(0.4)
return await super().__aenter__()
async def call_tool(self, name, args):
await asyncio.sleep(0.5)
return await super().call_tool(name, args)
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowBoth(url))
start = time.monotonic()
# 0.4s connect + 0.5s call vs a 0.6s budget: the call must inherit only
# the remaining ~0.2s, not a fresh full window.
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.6, scope = "chat")
assert "timed out" in out
assert time.monotonic() - start < 2.0
def test_close_narrowed_by_headers_spares_other_env(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
call_tool_sync(STDIO_URL, {"ENV_VAR": "b"}, "t", {}, scope = "chat")
# Two server rows can share a command with different envs; editing one
# must only close its own sessions.
close_stdio_sessions(STDIO_URL, None)
assert fake_clients[0].exited == 1
assert fake_clients[1].exited == 0
assert len(mcp_client._stdio_sessions) == 1
close_stdio_sessions(STDIO_URL) # headers omitted: any env for the command
assert fake_clients[1].exited == 1
assert mcp_client._stdio_sessions == {}
def test_close_stdio_sessions_by_url(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
call_tool_sync("npx other-server", None, "t", {}, scope = "chat")
key = mcp_client._session_key(STDIO_URL, None, "chat")
close_stdio_sessions(STDIO_URL)
assert fake_clients[0].exited == 1
assert fake_clients[1].exited == 0
assert len(mcp_client._stdio_sessions) == 1
assert key not in mcp_client._stdio_key_locks
def test_execute_tool_mcp_scope_is_per_thread(tmp_path, monkeypatch):
# session_id is the sandbox id and can be shared project-wide; the stdio
# session scope must also carry the per-conversation thread id.
from core.inference import tools as tools_mod
from storage import mcp_servers_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
monkeypatch.setattr(tools_mod, "stdio_mcp_enabled", lambda: True)
mcp_servers_db.create_server(id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True)
scopes: list = []
def fake_call_tool_sync(**kwargs):
scopes.append(kwargs["scope"])
return "ok"
monkeypatch.setattr(tools_mod, "call_tool_sync", fake_call_tool_sync)
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-a")
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-b")
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "sess-only")
tools_mod.execute_tool("mcp__s1__t", {}, thread_id = "thread-a")
# Persist only with a thread_id; session_id alone stays one-shot (None) so a
# project-wide id can't leak state across conversations. Fields are tagged.
assert scopes == ["s=project-p1:t=thread-a", "s=project-p1:t=thread-b", None, "s=:t=thread-a"]
# IDs containing ":" must not collapse distinct conversations into one scope,
# and a session-only id must never collide with a thread-only id.
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "a:b", thread_id = "c")
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "a", thread_id = "b:c")
assert scopes[-2] != scopes[-1]
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "same")
tools_mod.execute_tool("mcp__s1__t", {}, thread_id = "same")
assert scopes[-2] != scopes[-1] # session-only "same" != thread-only "same"
def test_execute_tool_config_check_tracks_row(tmp_path, monkeypatch):
from core.inference import tools as tools_mod
from storage import mcp_servers_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
monkeypatch.setattr(tools_mod, "stdio_mcp_enabled", lambda: True)
mcp_servers_db.create_server(id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True)
captured: dict = {}
monkeypatch.setattr(tools_mod, "call_tool_sync", lambda **kw: captured.update(kw) or "ok")
tools_mod.execute_tool("mcp__s1__t", {})
check = captured["config_check"]
assert check() is True
mcp_servers_db.update_server("s1", {"url": "npx different-server"})
assert check() is False
def test_multi_block_result_flattens_through_session(fake_clients):
async def _rich_call(name, args):
return SimpleNamespace(
content = [
SimpleNamespace(type = "text", text = "### Page"),
SimpleNamespace(type = "text", text = "- Page URL: https://example.com/"),
],
is_error = False,
structured_content = None,
)
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
fake_clients[0].call_tool = _rich_call
out = call_tool_sync(STDIO_URL, None, "browser_snapshot", {}, scope = "chat")
assert out == "### Page\n- Page URL: https://example.com/"
def test_stdio_cache_trims_overshoot_after_burst(fake_clients, monkeypatch):
# A concurrent burst of distinct-scope calls can overshoot the cap while every
# session is busy (insert-time eviction only reclaims idle sessions). Once the
# calls finish, release-time trimming must bring the cache back within cap.
monkeypatch.setattr(mcp_client, "_STDIO_MAX_SESSIONS", 2)
def slow_client(
url,
headers,
use_oauth = False,
):
client = FakeClient(url)
client.call_delay = 0.5 # keep every session in-flight during the burst
return client
monkeypatch.setattr(mcp_client, "_client", slow_client)
errors: list = []
def worker(i: int):
try:
call_tool_sync(STDIO_URL, None, "t", {}, scope = f"chat-{i}")
except Exception as exc: # noqa: BLE001
errors.append(exc)
threads = [threading.Thread(target = worker, args = (i,)) for i in range(5)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(10.0)
assert not errors, errors
assert len(mcp_client._stdio_sessions) <= 2
def test_close_http_server_creates_no_stdio_tombstone(fake_clients):
# HTTP/SSE servers are never cached as stdio sessions, so closing one on
# update/delete must not accrue a close-generation entry (an unbounded leak).
before_cfg = len(mcp_client._stdio_cfg_close_gen)
before_url = len(mcp_client._stdio_url_close_gen)
for i in range(50):
close_stdio_sessions(f"https://mcp-{i}.example/mcp", {"K": str(i)})
close_stdio_sessions(f"https://mcp-{i}.example/mcp")
assert len(mcp_client._stdio_cfg_close_gen) == before_cfg
assert len(mcp_client._stdio_url_close_gen) == before_url
# a real stdio command still registers a generation (the guard is non-stdio only)
close_stdio_sessions(STDIO_URL, {"K": "v"})
assert len(mcp_client._stdio_cfg_close_gen) == before_cfg + 1

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

@ -4,6 +4,8 @@ import sys
import types
from types import SimpleNamespace
import pytest
class _DummyMetal:
@staticmethod
@ -185,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri
assert isinstance(backend._tokenizer, _DummyTokenizer)
def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch):
_install_fake_mlx(monkeypatch)
calls = []
_install_fake_fast_mlx(monkeypatch, calls)
from core.inference.mlx_inference import MLXInferenceBackend
group = SimpleNamespace(size = lambda: 2, rank = lambda: 0)
config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False)
for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")):
calls.clear()
assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group)
_, kwargs = calls.pop()
assert kwargs["text_only"] is False and kwargs[group_key] is group
calls.clear()
singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0)
assert MLXInferenceBackend().load_model(
config, parallel_mode = "tensor", distributed_group = singleton
)
assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1])
config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True)
with pytest.raises(ValueError, match = "LoRA adapter repos"):
MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group)
@pytest.mark.parametrize("accepts_backend", (True, False))
def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend):
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import _init_mlx_distributed
group = SimpleNamespace(rank = lambda: 1, size = lambda: 2)
calls = []
def _init(**kwargs):
calls.append(kwargs)
if kwargs and not accepts_backend:
raise TypeError("backend keyword unsupported")
return group
sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init)
monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345")
monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json")
assert _init_mlx_distributed() == (group, 1, 2)
assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}])
def test_worker_share_object_receives_distributed_payload(monkeypatch):
from core.inference import worker
shared_obj = {"type": "turn", "text": "hi"}
payload = worker._encode_share_object(shared_obj)
def _array(value):
val = value.item() if hasattr(value, "item") else value
return SimpleNamespace(
item = lambda: val,
tolist = lambda: list(val) if hasattr(val, "__iter__") else [val],
)
mlx_pkg = types.ModuleType("mlx")
mlx_core = types.ModuleType("mlx.core")
mlx_core.uint8 = "uint8"
mlx_core.array = _array
mlx_core.zeros = lambda *_a, **_k: _array([])
def _all_sum(value, group = None):
value = value.item() if hasattr(value, "item") else value
return _array(len(payload)) if value == 0 else _array(payload)
mlx_core.distributed = SimpleNamespace(all_sum = _all_sum)
mlx_pkg.core = mlx_core
monkeypatch.setitem(sys.modules, "mlx", mlx_pkg)
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
responses = []
worker._handle_share_object(
SimpleNamespace(
_distributed_group = object(),
_distributed_rank = 1,
_distributed_world_size = 2,
),
{"type": "share_object", "request_id": "rid", "object": None},
SimpleNamespace(put = responses.append),
)
response = responses[0]
assert response["object"] == shared_obj
def test_worker_share_object_oversize_notifies_peers(monkeypatch):
from core.inference import worker
calls = []
mlx_pkg = types.ModuleType("mlx")
mlx_core = types.ModuleType("mlx.core")
mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value)
mlx_core.eval = lambda value: value
mlx_core.distributed = SimpleNamespace(
all_sum = lambda value, group = None: calls.append(value.item()) or value
)
mlx_pkg.core = mlx_core
monkeypatch.setitem(sys.modules, "mlx", mlx_pkg)
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8)
responses = []
worker._handle_share_object(
SimpleNamespace(
_distributed_group = object(),
_distributed_rank = 0,
_distributed_world_size = 2,
),
{"type": "share_object", "request_id": "rid", "object": {"text": "too long"}},
SimpleNamespace(put = responses.append),
)
assert calls == [worker._SHARE_OBJECT_ERROR_SIZE]
assert responses[0]["type"] == "share_error"
# Regression: generate_chat_response must accept the four template kwargs
# (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route
# layer can forward UI toggles. The old signature raised
@ -208,6 +333,118 @@ def test_mlx_generate_chat_response_accepts_template_kwargs():
), f"{name!r} must default to None so existing callers stay valid"
def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch):
from core.inference.mlx_inference import MLXInferenceBackend
calls = {"generic": [], "model": [], "stream": []}
state = {"generic": "serialized", "model": "<image> model-aware"}
prompt_utils = SimpleNamespace(
MODEL_CONFIG = {"deepseek_vl_v2": object()},
apply_chat_template = lambda *_args, **kwargs: (
calls["model"].append(kwargs) or state["model"]
),
)
mlx_vlm = types.ModuleType("mlx_vlm")
mlx_vlm.prompt_utils = prompt_utils
mlx_vlm.stream_generate = lambda *_args, **kwargs: (
calls["stream"].append((_args, kwargs))
or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)])
)
monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm)
def generic(_target, _messages, **kwargs):
calls["generic"].append(kwargs)
if isinstance(state["generic"], Exception):
raise state["generic"]
if state["generic"] == "serialized":
return f"User: {_messages[0]['content']}"
return state["generic"]
monkeypatch.setattr(
"core.inference.chat_template_helpers.apply_chat_template_for_generation",
generic,
)
backend = MLXInferenceBackend()
backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"})
backend._processor = SimpleNamespace(tokenizer = SimpleNamespace())
args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None)
tools = [{"function": {"name": "search"}}]
assert list(backend._generate_vlm(*args)) == ["ok"]
assert calls["model"][0]["num_images"] == 1
assert calls["stream"][0][0][2] == "<image> model-aware"
with pytest.raises(RuntimeError, match = "dropping requested tools"):
list(backend._generate_vlm(*args, tools = tools))
with pytest.raises(RuntimeError, match = "dropping requested tools or reasoning"):
list(backend._generate_vlm(*args, enable_thinking = False))
backend._processor = SimpleNamespace(chat_template = "template")
state["generic"] = "<image> healthy generic"
assert list(backend._generate_vlm(*args, tools = tools, enable_thinking = False)) == ["ok"]
assert calls["generic"][-1]["enable_thinking"] is False
assert calls["stream"][-1][0][2] == "<image> healthy generic"
state["generic"] = "generic prompt"
text_messages = [{"role": "user", "content": "hello"}]
assert list(backend._generate_vlm(*((text_messages, None) + args[2:]), tools = tools)) == ["ok"]
assert calls["generic"][-1]["tools"] == tools
assert calls["stream"][-1][0][2] == "generic prompt"
two_images = [{"role": "user", "content": [{"type": "image"}, {"type": "image"}]}]
with pytest.raises(RuntimeError, match = "2 structured image item"):
list(backend._generate_vlm(*((two_images,) + args[1:]), tools = tools))
state["generic"] = "serialized"
tool_history = args[0] + [{"role": "assistant", "tool_calls": [{"id": "call-1"}]}]
with pytest.raises(RuntimeError, match = "tool-call history"):
list(backend._generate_vlm(*((tool_history,) + args[1:]), tools = tools))
state["generic"] = ValueError("generic rendering failed")
state["model"] = f"User: {args[0][0]['content']}"
with pytest.raises(ValueError, match = "generic rendering failed"):
list(backend._generate_vlm(*args))
def test_mlx_vlm_image_injection_reuses_media_aliases(monkeypatch):
from core.inference.mlx_inference import MLXInferenceBackend, _prompt_serializes_vlm_media
media = [{"type": "image"}]
quoted = [{"role": "user", "content": media}, {"role": "user", "content": f"Explain {media}"}]
assert _prompt_serializes_vlm_media(f"<image>\n{media[0]}", quoted[:1])
assert not _prompt_serializes_vlm_media(f"<image>\nExplain {media}", quoted)
assert _prompt_serializes_vlm_media(f"User: {media}\nExplain {media}", quoted)
quoted[1]["content"] = [{"type": "text", "text": f'Explain "this" {media}'}]
assert not _prompt_serializes_vlm_media(f'<image>\nExplain "this" {media}', quoted)
json_media = [{"type": "image_url"}]
json_repr = '{"type": "image_url"}'
assert _prompt_serializes_vlm_media(f"<image>\n{json_repr}", [{"content": json_media}])
assert not _prompt_serializes_vlm_media(
f"<image>\nExplain {json_repr}",
[{"content": json_media}, {"content": f"Explain {json_repr}"}],
)
backend = MLXInferenceBackend()
backend._model = object()
backend._is_vlm = True
captured = []
backend._generate_vlm = lambda messages, *_args, **_kwargs: (
captured.append(messages) or iter(())
)
messages = [{"role": "user", "content": [{"type": "image_url"}]}]
list(backend.generate_chat_response(messages, image = object()))
assert captured[0][0]["content"] == [{"type": "image_url"}]
def test_mlx_vlm_model_config_prefers_config_with_model_type():
from core.inference.mlx_inference import _mlx_vlm_model_config
# config present but missing model_type must fall back to _config
m = SimpleNamespace(config = {}, _config = {"model_type": "deepseek_vl_v2"})
assert _mlx_vlm_model_config(m) == ({"model_type": "deepseek_vl_v2"}, "deepseek_vl_v2")
# an object config whose model_type is None also falls back
m = SimpleNamespace(config = SimpleNamespace(model_type = None), _config = {"model_type": "qwen2_vl"})
assert _mlx_vlm_model_config(m)[1] == "qwen2_vl"
# a config that already carries a model_type is preferred and returned unchanged
assert _mlx_vlm_model_config(SimpleNamespace(config = {"model_type": "gemma3"})) == (
{"model_type": "gemma3"},
"gemma3",
)
def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
"""Mac text path must route through apply_chat_template_for_generation so
reasoning / tool kwargs reach the tokenizer."""

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