Merge origin/main into fix/web-fetch-binary-guard

Resolve import conflict in studio/backend/core/inference/tools.py by keeping
both codecs (web fetch binary guard) and fnmatch (chat tool permission levels).
This commit is contained in:
danielhanchen 2026-07-15 13:30:04 +00:00
commit 76de937a95
101 changed files with 17870 additions and 854 deletions

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,6 +362,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 \

View file

@ -485,7 +485,7 @@ jobs:
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
@ -501,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()
@ -513,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",
@ -669,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({
@ -740,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,
@ -748,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)"

View file

@ -471,11 +471,22 @@ jobs:
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(
@ -487,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 = {
@ -575,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,
@ -583,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
@ -616,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}")

View file

@ -677,7 +677,22 @@ jobs:
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(
@ -689,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 = {
@ -749,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,
@ -757,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"
@ -780,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"
@ -802,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}")

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. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind.
- `-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.

View file

@ -91,6 +91,7 @@ function Install-UnslothStudio {
if ($TauriMode) {
exit $Code
}
throw $Message
}
# ── Parse flags ──
@ -2627,8 +2628,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 {
@ -2648,8 +2649,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 ""
}
}

View file

@ -3266,8 +3266,8 @@ if [ "$_SKIP_AUTOSTART" != true ] && [ -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
@ -3288,7 +3288,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

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",

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()

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

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,

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."""
@ -409,14 +447,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 +628,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)

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.

View file

@ -8969,16 +8969,37 @@ class LlamaCppBackend:
disable_parallel_tool_use: bool = False,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
) -> Generator[dict, None, None]:
"""
Agentic loop: let the model call tools, execute them, and continue.
permission_mode: "ask" confirms every call (with confirm_tool_calls),
"auto" only pauses calls detected as potentially unsafe, "off" never
pauses (sandbox stays on), "full" is the same as bypass_permissions.
Unset/unknown behaves as "ask".
Yields dicts:
{"type": "status", "text": "Searching: ..."/"Reading: ..."} -- tool status updates
{"type": "content", "text": "token"} -- streamed content tokens (cumulative)
{"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative)
"""
from core.inference.tools import build_rag_autoinject, execute_tool
from core.inference.tools import (
build_rag_autoinject,
execute_tool,
is_always_safe_tool,
is_potentially_unsafe_tool_call,
)
# Normalize the mode: "full" and bypass_permissions are the same
# switch, whichever arrives first wins toward the permissive side.
# "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"
if not self.is_loaded:
raise RuntimeError("llama-server is not loaded")
@ -8986,8 +9007,14 @@ class LlamaCppBackend:
conversation = list(messages)
# Forced first-pass RAG so a doc question doesn't lose to web_search. Emits
# the same tool card + citations a real call would.
_auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope)
# the same tool card + citations a real call would. Skip it only when a
# retrieval call would actually prompt (ask mode); auto never gates the
# safe search_knowledge_base tool, so retrieval must still run there.
# off never prompts either, so it also keeps first-pass retrieval.
_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
@ -9357,8 +9384,16 @@ class LlamaCppBackend:
in provisional_started_tool_calls.values()
)
# Later parallel cards only reconcile when parallel use is enabled.
# In auto mode an always-safe tool (render_html) never
# prompts, so it must stream its early card too; mirror
# that here instead of gating on the raw confirm flag.
_confirm_gated = (
confirm_tool_calls and not bypass_permissions
confirm_tool_calls
and not bypass_permissions
and not (
permission_mode == "auto"
and is_always_safe_tool(current_name)
)
)
# Keep small-argument tools on the normal path.
_args_len = len(
@ -9925,7 +9960,18 @@ class LlamaCppBackend:
# Bypass wins over the confirm gate at the loop level too,
# so a direct internal caller with both flags never prompts.
needs_confirm = bool(confirm_tool_calls) and not bypass_permissions
# In "auto" mode only calls detected as potentially unsafe
# pause; read-only calls run straight through. "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":
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

View file

@ -5,15 +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)
@ -422,9 +527,7 @@ class MLXInferenceBackend:
{"type": "text", "text": content},
]
elif isinstance(content, list):
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
@ -632,17 +735,87 @@ 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
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

View file

@ -174,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,
@ -210,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()
@ -252,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."""
@ -882,6 +919,13 @@ class InferenceOrchestrator:
# 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
@ -935,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)
@ -1030,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
@ -1061,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
@ -1297,6 +1372,7 @@ class InferenceOrchestrator:
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,
@ -1364,6 +1440,7 @@ class InferenceOrchestrator:
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,
permission_mode = permission_mode,
)
def generate_with_adapter_control(

View file

@ -428,6 +428,7 @@ def run_safetensors_tool_loop(
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.
@ -453,10 +454,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
@ -539,7 +557,16 @@ def run_safetensors_tool_loop(
# 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 = ""
@ -1056,8 +1083,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()

File diff suppressed because it is too large Load diff

View file

@ -291,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

View file

@ -931,7 +931,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"

View file

@ -39,6 +39,27 @@ from utils.paths import outputs_root
logger = get_logger(__name__)
def _env_int(name: str, default: int) -> int:
try:
raw = (os.environ.get(name) or "").strip()
return int(raw) if raw else default
except ValueError:
return default
# Stop-watchdog escalation timeouts. Primary trigger: a short grace once "complete"
# (save done). Absolute cap is a backstop: long for save=True so a slow save is never
# killed mid-write, shorter for a cancel that has nothing to save.
_STOP_GRACE_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_GRACE_S", 15)
_STOP_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_TIMEOUT_S", 600)
_CANCEL_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_CANCEL_TIMEOUT_S", 120)
# Watchdog DB finalize: a few short retries so a transient SQLite lock doesn't lose the
# terminal state, since the watchdog is the sole finalizer once _proc is dropped.
_DB_FINALIZE_RETRIES = 3
_DB_FINALIZE_RETRY_S = 0.5
_pyplot = None
_pyplot_failed = False
@ -741,6 +762,13 @@ class TrainingBackend:
self._pump_running: bool = False
self._lock = threading.Lock()
# Stop watchdog: after a stop is requested, escalates to force_terminate()
# if the worker does not exit on its own within a bounded time. The watched
# proc is tracked so a new run always gets its own watcher.
self._stop_watchdog: Optional[threading.Thread] = None
self._stop_watchdog_proc: Optional[mp.Process] = None
self._complete_seen = threading.Event()
# Progress state (updated by pump thread from subprocess events)
self._progress = TrainingProgress()
self._should_stop = False
@ -765,6 +793,7 @@ class TrainingBackend:
self._metric_buffer: list[dict] = []
self._run_finalized: bool = False
self._db_run_created: bool = False
self._db_create_in_progress: bool = False
self._db_total_steps_set: bool = False
self._db_config: Optional[dict] = None
self._db_started_at: Optional[str] = None
@ -852,91 +881,125 @@ class TrainingBackend:
else:
defer_auto_selection = True
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
# Handshake with the sidecar install route: mark the spawn in progress BEFORE rechecking
# the reservation, so either this recheck aborts, or the install's is_training_active()
# sees this flag (or the recorded proc) and refuses.
from utils.transformers_version import sidecar_swap_in_progress
if defer_auto_selection:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs)
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
self._spawn_in_progress = True
if sidecar_swap_in_progress():
self._spawn_in_progress = False
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
"retry when it completes."
)
# Any exception between the handshake above and the flag reset below would
# otherwise leave _spawn_in_progress latched, wedging is_training_active
# (and the install route) until restart.
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory. Runs AFTER the handshake
# so a lost race to an install can't tear down chat/export for a training run that never spawns.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
if defer_auto_selection:
try:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
None, **gpu_selection_kwargs
)
except Exception:
# Flag is already set; a failed GPU selection must not leave is_training_active stuck True.
self._spawn_in_progress = False
raise
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
self._spawn_in_progress = False
return False
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
self.grad_norm_history.clear()
self.grad_norm_step_history.clear()
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_create_in_progress = False # a stale watchdog create can't block this run
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
self._spawn_in_progress = False
return True
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
return False
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
self.grad_norm_history.clear()
self.grad_norm_step_history.clear()
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
return True
self._spawn_in_progress = False
raise
def stop_training(self, save: bool = True) -> bool:
"""Send stop signal to the training subprocess."""
@ -953,15 +1016,212 @@ class TrainingBackend:
self._progress.status_message = (
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
)
# Guarantee the run finalizes even if the worker wedges after saving.
self._start_stop_watchdog(cancel = not save)
return True
def force_terminate(self) -> None:
"""Force-kill the training subprocess so state can be reset immediately."""
def _start_stop_watchdog(self, cancel: bool) -> None:
"""Start a daemon that force-terminates the worker if a requested stop does not
exit on its own. No-op if no worker is alive or a live watchdog already watches
this proc (a stale watchdog on an old proc never blocks a new run's watcher)."""
with self._lock:
if self._proc is not None and self._proc.is_alive():
logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid)
self._proc.terminate()
proc = self._proc
if proc is None or not proc.is_alive():
return
if (
self._stop_watchdog is not None
and self._stop_watchdog.is_alive()
and self._stop_watchdog_proc is proc
):
return
watchdog = threading.Thread(
target = self._stop_watchdog_loop,
args = (proc, cancel, self.current_job_id),
name = f"stop-watchdog-{self.current_job_id or 'unknown'}",
daemon = True,
)
self._stop_watchdog = watchdog
self._stop_watchdog_proc = proc
watchdog.start()
def _stop_watchdog_loop(
self,
target_proc: "mp.Process",
cancel: bool,
watched_job_id: Optional[str] = None,
) -> None:
"""Escalate a stuck stop to force_terminate(): grace after "complete", else the
absolute backstop (see the module timeouts). No-ops on a clean exit; exits
silently if a new run replaces the worker."""
started = time.monotonic()
complete_at: Optional[float] = None
reason = ""
while True:
with self._lock:
superseded = self._proc is not target_proc
# A later cancel has nothing to save, so tighten an in-flight save
# watchdog to the shorter cancel cap.
cancelling = cancel or self._cancel_requested
if superseded or not target_proc.is_alive():
return
now = time.monotonic()
abs_timeout = _CANCEL_TIMEOUT_S if cancelling else _STOP_TIMEOUT_S
if complete_at is None and self._complete_seen.is_set():
complete_at = now
if complete_at is not None and now - complete_at >= _STOP_GRACE_S:
reason = "worker still alive after save"
break
if now - started >= abs_timeout:
reason = "worker did not exit within the absolute timeout"
break
time.sleep(0.5)
with self._lock:
superseded = self._proc is not target_proc
if superseded or not target_proc.is_alive():
return
if complete_at is None:
# Backstop fired pre-completion: a save may still be in progress.
logger.warning(
"Stop watchdog: absolute timeout with no completion signal; "
"force-terminating a possibly-mid-save worker: %s",
reason,
)
else:
logger.warning("Stop watchdog force-terminating stuck training worker: %s", reason)
# force_terminate can raise on a wedged child; finalize regardless.
try:
self.force_terminate(target_proc = target_proc)
except Exception:
logger.exception("Stop watchdog: force_terminate failed; finalizing anyway")
finally:
self._finalize_stopped_after_escalation(
target_proc = target_proc, watched_job_id = watched_job_id
)
def _finalize_stopped_after_escalation(
self,
target_proc: "Optional[mp.Process]" = None,
watched_job_id: Optional[str] = None,
) -> None:
"""Finalize parent state after a force-terminate so the UI leaves "Stopping..."
even if the worker is wedged in driver teardown; preserves output_dir so a saved
checkpoint is kept. No-ops if a new run already replaced the watched worker, so a
stale watchdog never marks a fresh run stopped or drops its handle.
Supersession is checked on both the watched proc and job id: start_training sets
current_job_id before it installs the new _proc, so a stale watchdog entering that
startup window still sees the old (dead) handle and is caught by the job-id guard.
The run's terminal DB state is recorded (create-if-needed + finish by captured id)
BEFORE _proc is dropped: a wedged worker still reports alive, so the pump never
reaches its own finalize and would bail on its _proc-is-None guard once the handle
is gone. While the handle is held is_training_active() stays true, so no new run can
start and current_job_id stays the watched run for the write. _proc is dropped last,
re-guarded on target_proc so a run that did replace the worker keeps its handle."""
with self._lock:
if target_proc is not None and self._proc is not target_proc:
return # a new run replaced the worker; never touch its state
if watched_job_id is not None and self.current_job_id != watched_job_id:
return # a new run is already starting up; leave its state alone
run_id = self.current_job_id # == watched_job_id
self._progress.is_training = False
self._progress.status_message = "Training stopped."
# Create the row if a start-time create failed (no-op otherwise; skips when the pump
# is mid-create, in which case its create-then-finalize records the run instead).
self._ensure_db_run_created()
with self._lock:
claim = (
bool(run_id)
and self.current_job_id == run_id
and self._db_run_created
and not self._run_finalized
)
batch: list = []
final_step = final_loss = duration = None
loss_history: list = []
output_dir = self._output_dir
if claim:
self._run_finalized = True # claim this run's finalize
batch = list(self._metric_buffer)
del self._metric_buffer[: len(batch)]
final_step = self._progress.step
final_loss = self._progress.loss
if final_loss is not None and not math.isfinite(final_loss):
final_loss = None
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
if claim:
self._finish_stopped_run(
run_id, output_dir, batch, final_step, final_loss, duration, loss_history
)
with self._lock:
if target_proc is None or self._proc is target_proc:
self._proc = None # drop only our handle, never a run that replaced it
def _finish_stopped_run(
self,
run_id: str,
output_dir: Optional[str],
batch: list,
final_step: Optional[int],
final_loss: Optional[float],
duration: Optional[float],
loss_history: list,
) -> None:
"""Record a force-stopped run finished by its captured id, from state snapshotted
under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE,
so a concurrent pump finalize of the same run is harmless and a different current run
is never touched. The watchdog is the sole finalizer once _proc is dropped, so a
transient DB error (e.g. a SQLite lock) is retried a few times; on final failure the
finalize is unclaimed (only if the run is still current) so the row is not left
claimed-but-unfinalized."""
for attempt in range(_DB_FINALIZE_RETRIES):
try:
from storage.studio_db import finish_run, insert_metrics_batch
from utils.downsample import downsample
if batch:
insert_metrics_batch(run_id, batch)
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
status = "stopped",
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = None,
)
return
except Exception:
if attempt + 1 < _DB_FINALIZE_RETRIES:
time.sleep(_DB_FINALIZE_RETRY_S)
continue
logger.warning(
"Failed to finalize stopped run %s in DB after %d attempts",
run_id,
_DB_FINALIZE_RETRIES,
exc_info = True,
)
with self._lock:
# Only if still current; a new run's finalize state is never touched.
if self.current_job_id == run_id:
self._run_finalized = False
def force_terminate(self, target_proc: "Optional[mp.Process]" = None) -> None:
"""Force-kill the training subprocess so state can be reset immediately. With
``target_proc``, terminate only that handle and no-op if a new run has replaced
it, so the watchdog can never kill a fresh worker."""
with self._lock:
proc = self._proc
if target_proc is not None and proc is not target_proc:
return # superseded by a new run; do not touch the new worker
if proc is not None and proc.is_alive():
logger.info("Force-terminating training subprocess (pid=%s)", proc.pid)
proc.terminate()
cancelled = self._cancel_requested
output_dir = self._output_dir
@ -1038,50 +1298,84 @@ class TrainingBackend:
from .worker import run_training_process
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
# This run is active, so an install request 409s rather than proceeds: a reservation seen here
# is transient (an aborting install or short lazy repair). Wait it out instead of stranding the
# stalled run; only a wedged reservation fails the respawn.
from utils.transformers_version import sidecar_swap_in_progress
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
self._spawn_in_progress = True
_swap_wait_deadline = time.time() + 120
while sidecar_swap_in_progress() and time.time() < _swap_wait_deadline:
time.sleep(1)
if sidecar_swap_in_progress():
# Raising here would land in the pump's broad finalization catch and
# strand the run in a training state with no worker: finalize it as a
# failure explicitly instead.
self._spawn_in_progress = False
msg = (
"A transformers installation is replacing the latest sidecar; "
"cannot respawn the training worker."
)
logger.error(msg)
with self._lock:
self._progress.is_training = False
self._progress.error = msg
self._ensure_db_run_created()
self._finalize_run_in_db(status = "error", error_message = msg)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
# Reset the handshake flag on any unexpected failure past this point, so a
# crashed respawn cannot wedge is_training_active until restart.
try:
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
self._spawn_in_progress = False
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._spawn_in_progress = False
self._pump_thread = new_pump
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
except Exception:
self._spawn_in_progress = False
raise
def _ensure_pump_alive(self) -> bool:
"""Restart the event pump if it crashed, even after the worker exited.
@ -1114,6 +1408,10 @@ class TrainingBackend:
def is_training_active(self) -> bool:
"""Check if training is currently active."""
# A spawn past its sidecar-swap recheck counts as active even before _proc is recorded,
# so an install cannot slip in mid-spawn.
if getattr(self, "_spawn_in_progress", False):
return True
# Self-heal a crashed pump first: a dead pump must never leave the worker
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
@ -1468,6 +1766,8 @@ class TrainingBackend:
"training cancelled",
"training stopped",
}
# Save is done by now; let the stop watchdog start its grace timer.
self._complete_seen.set()
self._progress.is_training = False
self._progress.is_completed = not stopped
self._output_dir = event.get("output_dir")
@ -1532,90 +1832,135 @@ class TrainingBackend:
self._finalize_run_in_db(**db_action_kwargs)
def _ensure_db_run_created(self) -> None:
"""Create the DB row if it doesn't exist yet. Called outside the lock."""
if self._db_run_created or not self.current_job_id or not self._db_config:
return
"""Create the DB row if it doesn't exist yet. An in-progress flag lets only one
caller create at a time, and ``_db_run_created`` is published only after
``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a
not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running)."""
with self._lock:
if (
self._db_run_created
or self._db_create_in_progress
or not self.current_job_id
or not self._db_config
):
return
self._db_create_in_progress = True # only one caller creates
job_id = self.current_job_id
db_config = self._db_config
started_at = self._db_started_at or datetime.now(timezone.utc).isoformat()
total_steps = self._progress.total_steps or None
created = False
try:
from storage.studio_db import create_run
dataset_name = (
self._db_config.get("hf_dataset")
or next(iter(self._db_config.get("local_datasets") or []), None)
or _s3_dataset_name(self._db_config.get("s3_dataset"))
db_config.get("hf_dataset")
or next(iter(db_config.get("local_datasets") or []), None)
or _s3_dataset_name(db_config.get("s3_dataset"))
or "unknown"
)
create_run(
id = self.current_job_id,
model_name = self._db_config["model_name"],
id = job_id,
model_name = db_config["model_name"],
dataset_name = dataset_name,
config_json = _json.dumps(self._db_config),
started_at = self._db_started_at or datetime.now(timezone.utc).isoformat(),
total_steps = self._progress.total_steps or None,
config_json = _json.dumps(db_config),
started_at = started_at,
total_steps = total_steps,
)
self._db_run_created = True
created = True
except Exception:
logger.warning("Failed to create DB run record for early failure", exc_info = True)
finally:
with self._lock:
# Publish the flags only if this is still the current run. A killed worker
# lets a new /start proceed mid-create, and these flags are backend-wide, so
# a stale create for the captured job must not satisfy the new run's DB state
# (the row was still created by id; the new run owns/creates its own row).
if self.current_job_id == job_id:
if created:
self._db_run_created = True # publish only after the insert commits
self._db_create_in_progress = False
def _finalize_run_in_db(
self,
status: str,
error_message: Optional[str] = None,
output_dir: Optional[str] = None,
expected_job_id: Optional[str] = None,
) -> None:
"""Flush remaining metrics and mark a run as finished in the DB."""
if not self.current_job_id or not self._db_run_created or self._run_finalized:
return
self._flush_metrics_to_db()
"""Flush remaining metrics and mark a run finished in the DB. Claims the finalize
under the lock so the watchdog and pump can't double-finalize, and no-ops when
``expected_job_id`` no longer matches (a new run took over). The run id and final
progress are snapshotted under the lock and threaded through the flush/finish calls,
so a new run racing between this claim and the DB writes can't be flushed or marked
stopped under the old run's finalize."""
with self._lock:
if expected_job_id is not None and self.current_job_id != expected_job_id:
return
if not self.current_job_id or not self._db_run_created or self._run_finalized:
return
self._run_finalized = True
run_id = self.current_job_id
final_step = self._progress.step
final_loss = self._progress.loss
if final_loss is not None and not math.isfinite(final_loss):
final_loss = None
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
self._flush_metrics_to_db(run_id = run_id)
try:
from storage.studio_db import finish_run
from utils.downsample import downsample
sparkline = downsample(self.loss_history, 50)
sparkline = downsample(loss_history, 50)
finish_run(
id = self.current_job_id,
id = run_id,
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = self._progress.step,
final_loss = self._progress.loss
if (self._progress.loss is not None and math.isfinite(self._progress.loss))
else None,
duration_seconds = self._progress.elapsed_seconds,
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = error_message,
)
self._run_finalized = True
except Exception:
with self._lock:
self._run_finalized = False # unclaim so a later flush can retry
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
def _flush_metrics_to_db(self) -> None:
"""Flush buffered metrics to the database and update live progress."""
if not self._metric_buffer or not self.current_job_id or not self._db_run_created:
return
# Cap buffer to bound memory growth.
if len(self._metric_buffer) > 500:
logger.warning(
"Metric buffer exceeded 500 entries (%d) — trimming oldest",
len(self._metric_buffer),
)
self._metric_buffer = self._metric_buffer[-500:]
# Snapshot before insert so metrics arriving during the write survive.
batch = list(self._metric_buffer)
def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None:
"""Flush buffered metrics to the DB and update live progress. The target run id,
metric batch, and progress snapshot are all taken under the lock, so a concurrent
flush can't double-remove metrics and a racing new run can't redirect the write to
a different job. A finalizer passes ``run_id`` to pin the target to its captured run."""
with self._lock:
target = run_id if run_id is not None else self.current_job_id
if not self._metric_buffer or not target or not self._db_run_created:
return
# Cap buffer to bound memory growth.
if len(self._metric_buffer) > 500:
logger.warning(
"Metric buffer exceeded 500 entries (%d) — trimming oldest",
len(self._metric_buffer),
)
del self._metric_buffer[:-500]
# Claim the batch under the lock so a concurrent flush can't re-remove it.
batch = list(self._metric_buffer)
del self._metric_buffer[: len(batch)]
step = self._progress.step
loss = self._progress.loss
if loss is not None and not math.isfinite(loss):
loss = None
duration = self._progress.elapsed_seconds
try:
from storage.studio_db import insert_metrics_batch, update_run_progress
insert_metrics_batch(self.current_job_id, batch)
del self._metric_buffer[: len(batch)]
update_run_progress(
id = self.current_job_id,
step = self._progress.step,
loss = self._progress.loss
if (self._progress.loss is not None and math.isfinite(self._progress.loss))
else None,
duration_seconds = self._progress.elapsed_seconds,
)
insert_metrics_batch(target, batch)
update_run_progress(id = target, step = step, loss = loss, duration_seconds = duration)
except Exception:
# Leave buffer intact for retry on next flush
# Re-queue the claimed batch at the front so it retries on the next flush.
with self._lock:
self._metric_buffer[:0] = batch
logger.warning("Failed to flush metrics to DB", exc_info = True)
@staticmethod

View file

@ -3019,11 +3019,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),

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:

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")

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,
@ -228,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)

View file

@ -234,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
@ -559,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"
@ -571,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",
@ -1363,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.
@ -1398,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():
@ -1410,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

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)",
)

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.",
@ -1040,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
@ -1695,6 +1834,10 @@ class AnthropicMessagesRequest(BaseModel):
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).",
@ -1736,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 ────────────────────────────────────────────

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:

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,

View file

@ -1693,6 +1693,9 @@ from models.inference import (
CompletionUsage,
ValidateModelRequest,
ValidateModelResponse,
TransformersUpgradeInfo,
InstallLatestTransformersRequest,
InstallLatestTransformersResponse,
TextContentPart,
ImageContentPart,
ImageUrl,
@ -2061,6 +2064,59 @@ def _explicit_studio_tool_loop_requested(payload) -> bool:
return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled))
def _permission_mode_confirm(payload) -> bool:
"""Effective confirm-gate intent for Studio's own local tool loop.
Honors the documented default that an unset permission_mode behaves as
"ask". An explicit confirm_tool_calls (True or False) wins; explicit
ask/auto always engage the gate (a non-streaming one is then rejected, since
it cannot prompt); off/full never prompt. An unset mode defaults to ask, but
that is only realizable on a streaming request, so a non-streaming unset
request keeps the legacy run-without-gate behavior instead of 400ing. Used
at the pre-switch guard and the per-backend tool paths so a forced tool loop
(CLI --enable-tools) with the default mode still gates streaming requests.
"""
if payload.confirm_tool_calls is not None:
return bool(payload.confirm_tool_calls)
mode = getattr(payload, "permission_mode", None)
if mode in ("ask", "auto"):
return True
if mode in ("off", "full"):
return False
return bool(getattr(payload, "stream", False))
def _confirm_gate_needs_stream(payload) -> bool:
"""Whether Studio's local tool-loop confirm gate still requires stream=true.
The gate can only prompt while streaming, so a non-streaming request that will
prompt must 400 up front. auto ("Approve for me") only prompts for a call the
classifier flags, so an auto request whose confirm is derived from the mode
(not an explicit confirm_tool_calls=true) and whose selectable tools are all
always-safe (web_search / RAG) never prompts and needs no stream. ask,
an explicit confirm flag, MCP tools, and an unrestricted or unsafe selection
still require streaming.
"""
if not _permission_mode_confirm(payload):
return False
if getattr(payload, "permission_mode", None) != "auto":
return True
if payload.confirm_tool_calls is True:
return True
if getattr(payload, "mcp_enabled", False):
return True
enabled = getattr(payload, "enabled_tools", None)
if enabled is None:
return True # omitted enabled_tools resolves to ALL tools (incl. terminal/python)
if not enabled:
# An explicit empty selection runs no built-in tool (_select_request_tools
# skips the loop), so there is nothing to prompt and no stream is needed.
return False
from core.inference.tools import is_always_safe_tool
return not all(is_always_safe_tool(t) for t in enabled)
# Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so
# is_disconnected() never fires. POST /inference/cancel looks up in-flight
# cancel_events here by cancel_id (per-run) or session_id / completion_id
@ -3836,11 +3892,25 @@ async def load_model(
GGUF models load via llama-server (llama.cpp) instead of Unsloth.
"""
# A sidecar install that has reserved the swap must not lose to a load that
# then gets unloaded by the pre-swap teardown. Rechecked under the gate: an
# install can reserve while this request queues on the gate, so the pre-gate
# check alone is only a fast path.
from core.inference.llama_keepwarm import inference_lifecycle_gate
from utils.transformers_version import sidecar_swap_in_progress
_swap_409 = HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
if sidecar_swap_in_progress():
raise _swap_409
# Hold the lifecycle gate across the load so idle auto-unload can't unload the
# model mid-load. Auto-switch calls _load_model_impl directly since it already
# holds this gate.
from core.inference.llama_keepwarm import inference_lifecycle_gate
async with inference_lifecycle_gate():
if sidecar_swap_in_progress():
raise _swap_409
return await _load_model_impl(request, fastapi_request, current_subject)
@ -4037,6 +4107,17 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
f"Resolved load_in_4bit={effective_load_in_4bit} for '{model_log_label}' "
f"from adapter_config.json / base model (requested {request.load_in_4bit})"
)
# Latest-sidecar models load 16-bit (worker refuses bnb 4-bit); size the guard
# to match. Off-loop: tier resolution reads configs.
if effective_load_in_4bit and not config.is_gguf:
from utils.transformers_version import latest_tier_active_for
if await asyncio.to_thread(latest_tier_active_for, config.identifier, request.hf_token):
effective_load_in_4bit = False
logger.info(
f"Latest-transformers sidecar active for '{model_log_label}' - "
"sizing and loading in 16-bit (4-bit is disabled for brand-new "
"architectures)"
)
# Refuse a load that would OOM active training, before the unload step below
# frees the resident model. Off-loop: guard does sync nvidia-smi / HF work.
@ -4470,6 +4551,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e)
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Lost the spawn-time race to a sidecar install/repair: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
# Friendlier message for models Unsloth cannot load.
if native_grant_backed:
redacted_msg = redact_native_paths(str(e))
@ -4598,16 +4684,6 @@ async def validate_model(
detail = "gpu_ids is not supported for GGUF models yet.",
)
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
)
# Both checks cover the [adapter, base] set (matching the scan route and workers):
# either repo can ship auto_map code or a poisoned pickle.
@ -4624,16 +4700,69 @@ async def validate_model(
security_targets = list(dict.fromkeys(security_targets))
is_gguf = getattr(config, "is_gguf", False)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
# run the remote-code/security preflight for non-GGUF loads.
# Does a newer transformers ship this model_type? Static overlay first, cached
# PyPI/main snapshot only for unknown types. Never fails validation; run before
# the training guard so an installable upgrade sizes as 16-bit.
transformers_upgrade: Optional[TransformersUpgradeInfo] = None
if not is_gguf:
from utils.transformers_latest import check_upgrade_for_model
# Cover [adapter, base]: the worker activates transformers for the base model.
for _target in security_targets:
_upgrade = await asyncio.to_thread(
check_upgrade_for_model, _target, request.hf_token
)
if _upgrade is not None:
transformers_upgrade = TransformersUpgradeInfo(**_upgrade)
break
# Whether the model can load on the CURRENT transformers through its own remote
# code (auto_map, or the YAML trust default). Computed before the 16-bit flip
# because a model with this fallback still loads 4-bit without the offered install,
# exactly as /load does.
requires_trust_remote_code = False
requires_security_review = False
if not is_gguf:
requires_trust_remote_code = any(
_requires_trust_remote_code_for_model(_t, request.hf_token)
for _t in security_targets
)
# Mirror /load's latest-sidecar 16-bit flip so the guard sizes it the same way. An
# ALREADY-ACTIVE latest sidecar always forces 16-bit (the worker will). A merely
# OFFERED (not yet installed) upgrade forces 16-bit only when the model has NO
# custom-code fallback: with auto_map it still loads 4-bit on the current
# transformers (as /load does without a successful install), and the install route
# refuses while training is active, so sizing 16-bit here would 409 the only viable
# 4-bit path. /load re-sizes 16-bit after a successful install and re-guards there.
if effective_load_in_4bit and not is_gguf:
from utils.transformers_version import latest_tier_active_for
_install_only_upgrade = (
transformers_upgrade is not None
and transformers_upgrade.supported_in_pypi
and transformers_upgrade.pypi_version
and not requires_trust_remote_code
)
if _install_only_upgrade or await asyncio.to_thread(
latest_tier_active_for, config.identifier, request.hf_token
):
effective_load_in_4bit = False
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
# run the security preflight for non-GGUF loads (requires_trust_remote_code was
# already resolved above for the sizing flip).
requires_security_review = False
if not is_gguf:
requires_security_review = any(
_requires_security_review_for_model(_t, request.hf_token) for _t in security_targets
)
@ -4676,6 +4805,8 @@ async def validate_model(
requires_trust_remote_code = requires_trust_remote_code,
requires_security_review = requires_security_review,
context_length = context_length,
requires_transformers_upgrade = transformers_upgrade is not None,
transformers_upgrade = transformers_upgrade,
)
except HTTPException:
@ -4720,6 +4851,217 @@ async def validate_model(
)
# studio_router only: admin action, kept off the OpenAI-compatible /v1 mount.
@studio_router.post(
"/install-latest-transformers", response_model = InstallLatestTransformersResponse
)
async def install_latest_transformers_route(
request: InstallLatestTransformersRequest, current_subject: str = Depends(get_current_subject)
):
"""
Consented install of the latest transformers release into the persistent
.venv_t5_latest sidecar.
Called after the user confirms the transformers-upgrade dialog raised by /validate
(requires_transformers_upgrade). The requested version must match the current latest
PyPI release (re-verified server-side); the sidecar then participates in routing on
this and every future start. A pip install runs off-loop, so this can take a minute.
"""
from utils.transformers_latest import install_latest_transformers
from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap
# The install stage-and-swaps .venv_t5_latest in place; a live worker would
# lazy-import from the new version mid-run, mixing incompatible modules. Gate on
# worker LIVENESS not tier (no HF token here, so tier re-resolution is unreliable
# for gated repos): training and export are refused, the chat model unloaded.
# Reserve the swap FIRST, before any await: training/export starts check this
# reservation, so raising it after the gate wait would let a worker slip in.
if not try_begin_sidecar_swap():
raise HTTPException(
status_code = 409,
detail = "A transformers installation is already in progress.",
)
# Until the installer thread takes over, this coroutine owns the reservation
# and must release it on any early exit (the 409 refusals below).
owns_reservation = True
try:
from core.export import get_export_backend
from core.training import get_training_backend
if get_training_backend().is_training_active():
raise HTTPException(
status_code = 409,
detail = (
"A training run is active. Wait for it to finish before "
"installing a new transformers version."
),
)
_export = get_export_backend()
if _export.is_export_active():
raise HTTPException(
status_code = 409,
detail = (
"An export is running. Wait for it to finish before "
"installing a new transformers version."
),
)
# A loaded (idle) export checkpoint would be torn down by the pre-swap
# cleanup; if the swap then failed, that state would be silently lost
# with no rollback signal. Make the user unload it deliberately first.
if getattr(_export, "current_checkpoint", None):
raise HTTPException(
status_code = 409,
detail = (
"An export checkpoint is loaded. Unload it from the Export "
"page before installing a new transformers version."
),
)
# In-flight streams passed the middleware already, so the lifecycle gate can't
# protect them and the swap's unload would kill them mid-stream; mirror the
# auto-switch busy check. This route is not middleware-counted and pending
# requests stay blocked in the middleware, so neither is subtracted here.
from core.inference.llama_keepwarm import (
inference_lifecycle_gate,
note_model_unloaded,
other_inference_request_count,
)
if other_inference_request_count(current_request_counted = False, include_pending = False) > 0:
raise HTTPException(
status_code = 409,
detail = (
"Another inference request is in progress. Wait for it to "
"finish before installing a new transformers version."
),
)
# Hold the lifecycle gate /load holds so no HF worker can start (or be mid-load
# with active_model_name unset) while the sidecar is swapped. Teardown runs via
# before_swap, only once the staged install succeeded: a failed pip/compat check
# must not leave the user with their model gone. GGUF stays loaded (llama-server
# never imports transformers).
backend = get_inference_backend()
export_backend = get_export_backend()
unloaded_chat = {"v": False}
def _unload_before_swap() -> None:
# Runs on the install thread, inside the gate held by _gated_install. Any
# failure raises so the previous sidecar stays untouched (a worker that did
# not tear down cleanly may still lazy-import from it). Export teardown runs
# FIRST so its failure aborts while the chat model is still loaded;
# cleanup_memory shuts the subprocess down even when its command fails, so
# judge by worker liveness, not its return value.
export_backend.cleanup_memory()
export_alive = getattr(export_backend, "is_worker_alive", None)
if callable(export_alive) and export_alive():
raise RuntimeError("Export worker still alive before the transformers swap")
active = getattr(backend, "active_model_name", None)
if active:
if not backend.unload_model(active):
# A failed unload still clears the orchestrator's model state,
# so the model is gone from the parent's view even though the
# swap aborts: report it so the client rolls back instead of
# pointing at an unloaded model.
if getattr(backend, "active_model_name", None) != active:
unloaded_chat["v"] = True
note_model_unloaded()
raise RuntimeError(f"Could not unload '{active}' before the transformers swap")
note_model_unloaded()
unloaded_chat["v"] = True
logger.info(
"Unloaded '%s' before swapping in transformers %s",
active,
request.version,
)
# A failed load can leave a live worker with no active model that
# still holds sidecar modules (and blocks the rename on Windows).
worker_alive = getattr(backend, "is_worker_alive", None)
if callable(worker_alive) and worker_alive():
# _shutdown_subprocess keeps the handle when the worker outlives SIGKILL,
# so both its False result and the liveness recheck catch a survivor
# rather than the recheck being fooled by a nulled handle.
stopped = backend._shutdown_subprocess()
if not stopped or worker_alive():
raise RuntimeError("Inference worker still alive before the transformers swap")
def _run_install() -> dict:
# Owns the reservation from here: releasing in the thread, not the route,
# keeps it held if the request is cancelled while the install still stages.
try:
return install_latest_transformers(request.version, _unload_before_swap, True)
finally:
end_sidecar_swap()
# Snapshot before waiting on the gate: a /load already holding it can
# complete meanwhile (including a same-model reload with new settings),
# and the installer must not unload a model whose successful LoadResponse
# the client is about to render. The generation counter catches reloads
# the name alone would miss.
active_before_gate = (
getattr(backend, "active_model_name", None),
getattr(backend, "load_generation", 0),
)
async def _gated_install() -> dict:
# Held by THIS task, not the request coroutine: a cancelled POST unwinding an
# `async with` here would drop the only guard /load honors mid-install.
async with inference_lifecycle_gate():
_active_now = (
getattr(backend, "active_model_name", None),
getattr(backend, "load_generation", 0),
)
if _active_now != active_before_gate:
end_sidecar_swap()
raise HTTPException(
status_code = 409,
detail = (
"A model load completed while the install was waiting. "
"Retry the install."
),
)
# Recheck under the gate: new streams bump their in-flight count while
# holding it, so once held nothing slips past (the pre-gate check is only
# a fast path and can be outlasted by a wait on a long /load).
if (
other_inference_request_count(
current_request_counted = False, include_pending = False
)
> 0
):
end_sidecar_swap()
raise HTTPException(
status_code = 409,
detail = (
"Another inference request is in progress. Wait for "
"it to finish before installing a new transformers "
"version."
),
)
return await asyncio.to_thread(_run_install)
install_task = asyncio.ensure_future(_gated_install())
owns_reservation = False
# shield: a cancelled request stops waiting, but the installer runs to
# completion (holding the gate) instead of being torn down mid-swap.
result = await asyncio.shield(install_task)
finally:
if owns_reservation:
end_sidecar_swap()
if not result["success"]:
if result.get("latest_version"):
# Structured failure so the dialog can update to the newer release
# and offer a retry that can actually succeed.
return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"])
if unloaded_chat["v"]:
# The chat model is already gone even though the swap failed; return a
# structured failure (not a bare 400) so the client can restore its
# model state instead of pointing at an unloaded model.
return InstallLatestTransformersResponse(**result, model_unloaded = True)
raise HTTPException(status_code = 400, detail = result["message"])
return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"])
@router.post("/unload", response_model = UnloadResponse)
async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)):
"""
@ -6322,25 +6664,52 @@ async def openai_chat_completions(
)
# Reject confirm-without-stream local tool requests before the switch: the
# local tool path requires stream=true for the confirm gate, so this shape
# is invalid and must not evict the resident model first. Mirror that path's
# enablement exactly (_effective_enable_tools honors a CLI --enable-tools
# policy hard-override; mcp_enabled opens the tool loop on its own but still
# defers to a CLI --disable-tools policy), or an mcp_enabled/policy-forced
# request would slip past this guard and only 400 after the swap.
from state.tool_policy import get_tool_policy as _get_confirm_tool_policy
# is invalid and must not evict the resident model first.
#
# Enter the local-loop arm exactly when the passthrough router below would
# run Studio's own tool loop. That gate is `_tools_on or _mcp_allowed`
# (see the use_tools block): _effective_enable_tools (which lets a
# process-wide --enable-tools policy force the loop on) plus mcp_enabled
# honoring --disable-tools, and tool_choice="none" disabling it unless the
# request explicitly asked. enabled_tools never enters loop entry (it only
# filters which tools run), so it is not a signal here.
#
# But a policy-forced loop must not steal client-tool passthrough: when the
# request did not explicitly ask for the loop (enable_tools/mcp) and carries
# client tools, the router forwards to the provider branch, so only treat it
# as the local loop when the request explicitly asked OR there is no client
# passthrough to defer to.
from state.tool_policy import get_tool_policy as _get_tool_policy_pre
_confirm_cli_policy = _get_confirm_tool_policy()
_cli_policy_pre = _get_tool_policy_pre()
_use_tools_intent = _effective_enable_tools(payload) or (
bool(payload.mcp_enabled) and _cli_policy_pre is not False
)
if payload.tool_choice == "none" and not _explicit_studio_tool_loop_requested(payload):
_use_tools_intent = False
_client_tool_passthrough = (
bool(payload.tools)
or bool(payload.openai_code_exec_container_id)
or bool(payload.anthropic_code_exec_container_id)
# A JSON-schema response_format is guided-decoding structured output the
# router forwards to the llama-server passthrough, not Studio's tool
# loop, so a --enable-tools policy must not 400 it as a local-confirm
# request under ask/auto.
or bool(_extract_response_format(payload))
)
# permission_mode only implies the confirm gate for that local loop.
# Client-tool passthrough forwards to the provider branch and the validator
# intentionally leaves confirm_tool_calls unset there, so only an explicit
# confirm_tool_calls=True should force the local-confirm rejection for it.
_studio_local_tool_loop = bool(_use_tools_intent) and (
_explicit_studio_tool_loop_requested(payload) or not _client_tool_passthrough
)
if (
payload.confirm_tool_calls
and not payload.bypass_permissions
not payload.bypass_permissions
and not payload.stream
and (
_effective_enable_tools(payload)
or (bool(payload.mcp_enabled) and _confirm_cli_policy is not False)
or bool(payload.enabled_tools)
or bool(payload.tools)
or bool(payload.openai_code_exec_container_id)
or bool(payload.anthropic_code_exec_container_id)
(_confirm_gate_needs_stream(payload) and _studio_local_tool_loop)
or (payload.confirm_tool_calls is True and _client_tool_passthrough)
)
):
raise HTTPException(
@ -6855,9 +7224,23 @@ async def openai_chat_completions(
use_tools = False
if use_tools:
# permission_mode ask/auto require the confirm gate for Studio's own
# tool loop. The request validator self-enables confirm only for
# request-level tool signals (enable_tools/enabled_tools/mcp_enabled);
# when a CLI policy (--enable-tools) forces the loop on without those,
# derive confirm here so the mode still gates the call (and a
# non-stream ask/auto request is rejected below rather than running
# unprompted). off/full never prompt, so they are excluded.
_effective_confirm = _permission_mode_confirm(payload)
# Bypass Permissions suppresses confirm, so the stream requirement
# (the gate needs streaming to prompt) no longer applies.
if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
# (the gate needs streaming to prompt) no longer applies. auto with an
# always-safe-only selection never prompts, so it needs no stream even
# though _effective_confirm stays true for the loop's per-call gate.
if (
_confirm_gate_needs_stream(payload)
and not payload.bypass_permissions
and not payload.stream
):
raise _reject(
400,
openai_error_body(
@ -6934,9 +7317,9 @@ async def openai_chat_completions(
disable_parallel_tool_use = payload.parallel_tool_calls is False,
# Bypass Permissions takes precedence over the confirm gate:
# never prompt while bypassing.
confirm_tool_calls = bool(payload.confirm_tool_calls)
and not bool(payload.bypass_permissions),
confirm_tool_calls = _effective_confirm and not bool(payload.bypass_permissions),
bypass_permissions = bool(payload.bypass_permissions),
permission_mode = payload.permission_mode,
)
_tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream"
@ -8150,9 +8533,20 @@ async def openai_chat_completions(
_sf_use_tools = False
if _sf_use_tools:
# permission_mode ask/auto require the confirm gate for Studio's own tool
# loop; when a CLI policy (--enable-tools) forces the loop on without a
# request-level tool signal, derive confirm here so the mode still gates
# the call (matching the GGUF path). off/full never prompt.
_sf_effective_confirm = _permission_mode_confirm(payload)
# Bypass Permissions suppresses confirm, so the stream requirement
# (the gate needs streaming to prompt) no longer applies.
if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
# (the gate needs streaming to prompt) no longer applies. auto with an
# always-safe-only selection never prompts, so it needs no stream even
# though _sf_effective_confirm stays true for the loop's per-call gate.
if (
_confirm_gate_needs_stream(payload)
and not payload.bypass_permissions
and not payload.stream
):
raise _reject(
400,
openai_error_body(
@ -8230,9 +8624,9 @@ async def openai_chat_completions(
rag_scope = payload.rag_scope,
# Bypass Permissions takes precedence over the confirm gate:
# never prompt while bypassing.
confirm_tool_calls = bool(payload.confirm_tool_calls)
and not bool(payload.bypass_permissions),
confirm_tool_calls = _sf_effective_confirm and not bool(payload.bypass_permissions),
bypass_permissions = bool(payload.bypass_permissions),
permission_mode = payload.permission_mode,
use_adapter = payload.use_adapter,
stats_holder = _sf_stats_holder,
)
@ -11248,6 +11642,13 @@ _STUDIO_ANTHROPIC_TOOL_ALIASES = {
"python": "python",
"terminal": "terminal",
}
# Server tools that never need a confirmation prompt (read-only / non code-
# executing; mirrors the unconditional-safe names in is_potentially_unsafe_tool_call).
# Any other selected tool (terminal, python, render_html) can require the gate
# this channel has no way to present, so an omitted permission_mode ("ask") only
# asks then. render_html is excluded because a networked canvas prompts in auto,
# and this channel invokes the loop without confirm; auto/ask reject, off/full run.
_ANTHROPIC_UNPROMPTED_SAFE_TOOLS = frozenset({"web_search", "search_knowledge_base"})
def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]:
@ -11508,6 +11909,53 @@ async def anthropic_messages(
),
)
# Reject an unsupported confirm-gated permission mode for Studio's own
# ("server") Anthropic tools before the switch, mirroring the malformed- and
# mixed-tool checks above. ask always wants a per-call pause this passthrough
# cannot offer, so it 400s whenever server tools are selected. auto only needs
# the gate for an unsafe call, so (like the omitted default) it runs for a
# safe-only selection (web_search/RAG) and 400s when a gate-needing tool is
# selected (local terminal/python, or render_html whose networked canvas
# prompts and cannot be gated on this channel). Rejecting must happen before the
# switch so an invalid request never evicts the resident model; it is
# determined from the requested tools alone (backend tool support is only known
# post-switch); an image request can never take the server-tool path, so it is
# excluded as in the server_tools gate below. off/full and an explicit
# confirm_tool_calls=False opt-out always pass.
_enable_pre = _effective_enable_tools(payload)
_server_tools_requested_pre = (
_enable_pre or (_enable_pre is None and bool(requested_studio_tools))
) and not _anthropic_request_has_image(payload)
if _server_tools_requested_pre:
from core.inference.tools import ALL_TOOLS as _ALL_TOOLS_PRE
_selected_pre = _select_anthropic_server_tools(
_ALL_TOOLS_PRE, requested_studio_tools, payload.enabled_tools
)
_perm_mode_pre = getattr(payload, "permission_mode", None)
_confirm_opt_out_pre = getattr(payload, "confirm_tool_calls", None) is False
_gated_tool_selected_pre = any(
tool["function"]["name"] not in _ANTHROPIC_UNPROMPTED_SAFE_TOOLS
for tool in _selected_pre
)
# An explicit confirm_tool_calls=False opts out of the gate entirely (it
# wins over the mode, mirroring _permission_mode_confirm and the GGUF path),
# so it never rejects -- not even under ask.
if not _confirm_opt_out_pre and (
_perm_mode_pre == "ask"
or (_perm_mode_pre in ("auto", None) and _gated_tool_selected_pre)
):
raise HTTPException(
status_code = 400,
detail = anthropic_error_body(
"permission_mode 'ask' has no confirmation channel for Anthropic "
"Messages server tools, and 'auto' (or the omitted default) cannot "
"gate a local 'terminal'/'python' tool here; set 'off' or 'full'.",
status = 400,
err_type = "invalid_request_error",
),
)
# require_vision rejects a swap to a text-only target before it runs, so an
# image request can't evict the resident vision model only to hit the vision
# guard (_normalize_anthropic_openai_images) below after the load.
@ -11707,6 +12155,10 @@ async def anthropic_messages(
)
from core.inference.tools import ALL_TOOLS
# ask/auto (and an omitted mode selecting a gate-needing terminal/python
# tool) were already rejected before the auto-switch above, so an invalid
# confirm-gated request never evicts the resident model; the selection
# here just picks the tools for the actual server-tool loop.
openai_tools = _select_anthropic_server_tools(
ALL_TOOLS,
requested_studio_tools,
@ -11762,6 +12214,7 @@ async def anthropic_messages(
rag_scope = getattr(payload, "rag_scope", None),
disable_parallel_tool_use = _disable_parallel,
bypass_permissions = bool(payload.bypass_permissions),
permission_mode = getattr(payload, "permission_mode", None),
)
if payload.stream:

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:

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

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.
@ -1202,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.")
@ -1315,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.
@ -1384,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
@ -1471,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",
@ -1487,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",
@ -1499,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(
@ -1561,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"
)
@ -1575,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)

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")

View file

@ -1739,7 +1739,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 +1763,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(

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

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):

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,

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,

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

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)

View file

@ -1826,6 +1826,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."""

View file

@ -333,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."""

View file

@ -706,6 +706,160 @@ class TestChatCompletionRequestToolFields:
assert entry["status"] == "completed"
assert monitor.active_count() == 0
def test_permission_mode_does_not_reject_client_tool_passthrough(self, monkeypatch):
# A non-streaming client-tool passthrough (client tools, no Studio tool
# loop) that also carries permission_mode "ask"/"auto" must reach the
# provider passthrough, not the confirm-without-stream guard: the
# validator leaves confirm_tool_calls unset for passthrough, and a bare
# permission_mode only gates Studio's own local tool loop. An explicit
# confirm_tool_calls=True still forces the local-confirm rejection.
# The pre-switch guard only runs when an automatic load may run, so force
# that predicate on to exercise it against a resident passthrough backend.
import routes.inference as inference_route
class _GGUFBackend:
is_loaded = True
model_identifier = "test-gguf"
supports_tools = False
supports_tool_passthrough = True
is_vision = False
_is_audio = False
context_length = 4096
base_url = "http://llama.permission-passthrough.test"
_request_reasoning_kwargs = lambda *_args, **_kwargs: None
def generate_chat_completion(self, **_kwargs):
raise AssertionError("client tools must use passthrough")
def generate_chat_completion_with_tools(self, **_kwargs):
raise AssertionError("Studio tool loop must stay disabled")
async def fake_passthrough(llama_backend, payload, model_name, **kwargs):
inference_route.api_monitor.finish(kwargs.get("monitor_id"))
return inference_route.JSONResponse({"ok": True, "model": model_name})
client_tools = [
{
"type": "function",
"function": {"name": "lookup", "parameters": {"type": "object"}},
}
]
def _setup(policy = None):
reset_tool_policy()
if policy is not None:
set_tool_policy(policy)
monkeypatch.setattr(inference_route, "_automatic_model_load_may_run", lambda: True)
monkeypatch.setattr(inference_route, "api_monitor", ApiMonitor(max_entries = 3))
monkeypatch.setattr(
inference_route, "_openai_passthrough_non_streaming", fake_passthrough
)
return self._v1_client(monkeypatch, _GGUFBackend())
# A process --enable-tools policy must not turn a client-tool passthrough
# into a Studio local loop, so a policy of None or True both keep the
# passthrough (the guard mirrors _explicit_studio_tool_loop_requested).
for policy in (None, True):
for mode in ("ask", "auto"):
client = _setup(policy)
resp = client.post(
"/v1/chat/completions",
json = {
"messages": [{"role": "user", "content": "use client tool"}],
"tools": client_tools,
"permission_mode": mode,
"stream": False,
},
)
assert resp.status_code == 200, resp.text
assert resp.json()["ok"] is True
# A JSON-schema response_format is guided-decoding passthrough, not a local
# tool loop, so a --enable-tools policy must not 400 a non-streaming ask/auto
# structured-output request under the confirm guard.
for mode in ("ask", "auto"):
client = _setup(True)
resp = client.post(
"/v1/chat/completions",
json = {
"messages": [{"role": "user", "content": "give me json"}],
"response_format": {
"type": "json_schema",
"json_schema": {"name": "s", "schema": {"type": "object"}},
},
"permission_mode": mode,
"stream": False,
},
)
assert resp.status_code == 200, resp.text
assert resp.json()["ok"] is True
# An explicit confirm_tool_calls=True with client tools and no stream is
# still a confirm-without-stream request and must be rejected up front.
client = _setup()
resp = client.post(
"/v1/chat/completions",
json = {
"messages": [{"role": "user", "content": "use client tool"}],
"tools": client_tools,
"confirm_tool_calls": True,
"stream": False,
},
)
assert resp.status_code == 400
assert "requires stream=true" in resp.json()["error"]["message"]
def test_permission_mode_policy_forced_local_loop_rejected_before_switch(self, monkeypatch):
# A process --enable-tools policy forces Studio's own tool loop on even
# when the request omits enable_tools and carries no client tools. A
# non-streaming ask/auto request is then confirm-gated with no stream to
# prompt on, so it must 400 at the pre-switch guard -- before
# _maybe_auto_switch_model runs -- rather than evicting the resident model
# and 400ing only at the per-backend check.
import routes.inference as inference_route
class _GGUFBackend:
is_loaded = True
model_identifier = "test-gguf"
supports_tools = True
supports_tool_passthrough = True
is_vision = False
_is_audio = False
context_length = 4096
base_url = "http://llama.policy-forced.test"
_request_reasoning_kwargs = lambda *_args, **_kwargs: None
switch_calls = []
async def _no_switch(*_args, **_kwargs):
switch_calls.append(1)
def _setup():
reset_tool_policy()
set_tool_policy(True)
monkeypatch.setattr(inference_route, "_automatic_model_load_may_run", lambda: True)
monkeypatch.setattr(inference_route, "api_monitor", ApiMonitor(max_entries = 3))
monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _no_switch)
return self._v1_client(monkeypatch, _GGUFBackend())
try:
for mode in ("ask", "auto"):
switch_calls.clear()
client = _setup()
resp = client.post(
"/v1/chat/completions",
json = {
"messages": [{"role": "user", "content": "hi"}],
"permission_mode": mode,
"stream": False,
},
)
assert resp.status_code == 400, resp.text
assert "requires stream=true" in resp.json()["error"]["message"]
assert switch_calls == [], "guard must reject before the auto-switch"
finally:
reset_tool_policy()
def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch):
# DiffusionGemma forces supports_tools off while passthrough stays
# available (#6851): enable_tools=True must not steal client tools

View file

@ -874,6 +874,35 @@ def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch):
assert o.models == {}
def test_load_model_aborts_when_old_worker_survives_shutdown(monkeypatch):
# A wedged worker that outlives terminate/kill makes _shutdown_subprocess return
# False. load_model must not spawn a second worker over it (double GPU allocation +
# the survivor's handle is lost); it aborts so the load can retry once it exits.
import types
from utils import transformers_version as tv
o = _bare_orchestrator()
o.active_model_name = "old"
o.models = {"old": {}}
o.loading_models = set()
monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False)
monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([0], "sel"))
monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None)
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
monkeypatch.setattr(o, "_cancel_generation", lambda: None)
monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: False) # survivor
monkeypatch.setattr(
o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn over a live survivor")
)
with pytest.raises(RuntimeError, match = "did not exit"):
o.load_model(types.SimpleNamespace(identifier = "new", gguf_variant = None))
# The except path cleared the loading marker and mirrors.
assert "new" not in o.loading_models
assert o.active_model_name is None
def test_load_model_proceeds_when_not_cancelled(monkeypatch):
# Guard against a false abort: an uncancelled load keeps its marker and spawns.
o = _bare_orchestrator()

View file

@ -0,0 +1,324 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Masked terminal password prompt (auth/terminal_prompt.py): reader echo and
editing, the change loop's validation/re-prompt behavior, and the pure
should-prompt gate. Drives the reader through a scripted fake getch, so no
tty (and no msvcrt on Linux) is needed."""
from __future__ import annotations
import io
import sys
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from auth import terminal_prompt as tp # noqa: E402
def _fake_getch(keys):
"""Scripted keystroke source: yields one item per _getch() call. Items may
be multi-char strings to simulate a paste burst arriving in one read."""
it = iter(keys)
def getch():
return next(it)
return getch
def _read(
monkeypatch,
keys,
prompt = "P: ",
):
monkeypatch.setattr(tp, "_getch", _fake_getch(keys))
out = io.StringIO()
value = tp._read_password(prompt, out = out)
return value, out.getvalue()
# ── _read_password ───────────────────────────────────────────────────
def test_reader_echoes_one_star_per_char(monkeypatch):
value, out = _read(monkeypatch, list("secret") + ["\r"])
assert value == "secret"
assert out.count("*") == 6
assert "secret" not in out
def test_reader_backspace_edits_and_erases_star(monkeypatch):
value, out = _read(monkeypatch, list("abc") + ["\x7f"] + list("d") + ["\n"])
assert value == "abd"
assert "\b \b" in out
# 4 stars were printed (a, b, c, d); one was erased.
assert out.count("*") == 4
def test_reader_backspace_on_empty_buffer_is_noop(monkeypatch):
value, out = _read(monkeypatch, ["\x08", "\x7f"] + list("x") + ["\r"])
assert value == "x"
assert "\b \b" not in out
def test_reader_paste_burst_delivers_all_chars(monkeypatch):
# A paste can arrive as one multi-char read; every char must count.
value, out = _read(monkeypatch, ["pasted-secret", "\r"])
assert value == "pasted-secret"
assert out.count("*") == len("pasted-secret")
def test_reader_unicode_password(monkeypatch):
value, _ = _read(monkeypatch, list("pässwörd✓") + ["\r"])
assert value == "pässwörd✓"
def test_reader_ignores_other_control_chars(monkeypatch):
value, _ = _read(monkeypatch, ["\t", "\x1b"] + list("ok") + ["\r"])
assert value == "ok"
def test_reader_ctrl_c_raises_keyboard_interrupt(monkeypatch):
monkeypatch.setattr(tp, "_getch", _fake_getch(list("ab") + ["\x03"]))
with pytest.raises(KeyboardInterrupt):
tp._read_password("P: ", out = io.StringIO())
def test_reader_ctrl_d_on_empty_raises_eof(monkeypatch):
monkeypatch.setattr(tp, "_getch", _fake_getch(["\x04"]))
with pytest.raises(EOFError):
tp._read_password("P: ", out = io.StringIO())
def test_reader_ctrl_d_mid_input_is_ignored(monkeypatch):
value, _ = _read(monkeypatch, list("ab") + ["\x04"] + list("c") + ["\r"])
assert value == "abc"
def test_reader_windows_key_prefix_is_ignored(monkeypatch):
# _getch_windows reports swallowed function-key sequences as "\x00".
value, _ = _read(monkeypatch, ["\x00"] + list("w") + ["\r"])
assert value == "w"
def test_reader_holds_raw_mode_once_for_whole_line(monkeypatch):
# Regression: cbreak/no-echo must be held for the ENTIRE line, not toggled
# per keystroke. Re-enabling echo between reads opens a window where a
# keystroke arriving in the gap echoes the password in cleartext. Assert the
# raw-mode context wraps the whole read exactly once and every keystroke is
# read while it is active.
events = []
class _SpyRawMode:
def __enter__(self):
events.append("enter")
return self
def __exit__(self, *exc):
events.append("exit")
return False
monkeypatch.setattr(tp, "_prompt_raw_mode", _SpyRawMode)
src = _fake_getch(list("s3cr3t!!") + ["\r"])
def _getch_recording():
assert events and events[-1] == "enter", "keystroke read outside raw mode"
return src()
monkeypatch.setattr(tp, "_getch", _getch_recording)
value = tp._read_password("P: ", out = io.StringIO())
assert value == "s3cr3t!!"
assert events == ["enter", "exit"]
# ── prompt_for_password_change ───────────────────────────────────────
def _run_loop(
monkeypatch,
keys,
*,
min_length = 8,
current = "bootstrap-pw",
):
monkeypatch.setattr(tp, "_getch", _fake_getch(keys))
out = io.StringIO()
applied = []
ok = tp.prompt_for_password_change(
min_length = min_length,
is_current_password = lambda pw: pw == current,
apply_change = applied.append,
out = out,
)
return ok, applied, out.getvalue()
def _keys(*lines):
keys = []
for line in lines:
keys.extend(list(line))
keys.append("\r")
return keys
def test_loop_success_applies_once(monkeypatch):
ok, applied, out = _run_loop(monkeypatch, _keys("new-password", "new-password"))
assert ok is True
assert applied == ["new-password"]
assert "Password updated" in out
assert "new-password" not in out
def test_loop_short_password_reprompts(monkeypatch):
ok, applied, out = _run_loop(monkeypatch, _keys("short", "long-enough-pw", "long-enough-pw"))
assert ok is True
assert applied == ["long-enough-pw"]
assert "at least 8 characters" in out
def test_loop_rejects_current_password(monkeypatch):
ok, applied, out = _run_loop(
monkeypatch, _keys("bootstrap-pw", "fresh-password", "fresh-password")
)
assert ok is True
assert applied == ["fresh-password"]
assert "must differ" in out
def test_loop_mismatch_reprompts_then_succeeds(monkeypatch):
ok, applied, out = _run_loop(
monkeypatch,
_keys("first-attempt", "typo-attempt", "second-attempt", "second-attempt"),
)
assert ok is True
assert applied == ["second-attempt"]
assert "do not match" in out
def test_loop_ctrl_c_aborts_without_applying(monkeypatch):
ok, applied, out = _run_loop(monkeypatch, list("ab") + ["\x03"])
assert ok is False
assert applied == []
assert "aborted" in out
def test_loop_eof_aborts_without_applying(monkeypatch):
ok, applied, out = _run_loop(monkeypatch, ["\x04"])
assert ok is False
assert applied == []
assert "aborted" in out
def test_loop_ctrl_c_at_confirmation_aborts(monkeypatch):
ok, applied, _ = _run_loop(monkeypatch, _keys("valid-password") + ["\x03"])
assert ok is False
assert applied == []
def test_loop_min_length_counts_code_points(monkeypatch):
# 8 unicode code points must pass a min_length of 8.
pw = "pässwörd"
assert len(pw) == 8
ok, applied, _ = _run_loop(monkeypatch, _keys(pw, pw))
assert ok is True
assert applied == [pw]
# ── should_prompt_password_change ────────────────────────────────────
@pytest.mark.parametrize(
"tunnel,requires,stdin_tty,stderr_tty,expected",
[
(True, True, True, True, True),
(False, True, True, True, False), # tunnel not starting (loopback no-op)
(True, False, True, True, False), # password already changed
(True, True, False, True, False), # piped stdin (headless)
(True, True, True, False, False), # redirected stderr
(False, False, False, False, False),
],
)
def test_should_prompt_matrix(tunnel, requires, stdin_tty, stderr_tty, expected):
assert (
tp.should_prompt_password_change(
tunnel_will_start = tunnel,
requires_change = requires,
stdin_isatty = stdin_tty,
stderr_isatty = stderr_tty,
)
is expected
)
def test_stream_eof_aborts_instead_of_submitting(monkeypatch):
# A dead stream ("" from _getch, e.g. a closed pty) must abort the line,
# never silently submit the partial password typed so far.
import io
err = io.StringIO()
monkeypatch.setattr(tp, "_getch", _fake_getch(list("abc") + [""]))
with pytest.raises(EOFError):
tp._read_password("New password: ", out = err)
# ── resolve_supplied_password: non-interactive --password / env / stdin ──
def test_resolve_supplied_password_literal_value_and_note(monkeypatch):
import io
monkeypatch.delenv(tp.SUPPLIED_PASSWORD_ENV, raising = False)
out = io.StringIO()
assert tp.resolve_supplied_password("hunter2pw", out = out) == "hunter2pw"
# A literal value warns that it is visible in the process list / history.
assert "process list" in out.getvalue()
def test_resolve_supplied_password_stdin(monkeypatch):
import io
monkeypatch.delenv(tp.SUPPLIED_PASSWORD_ENV, raising = False)
monkeypatch.setattr(sys, "stdin", io.StringIO("from-stdin-pw\n"))
assert tp.resolve_supplied_password("-") == "from-stdin-pw"
def test_resolve_supplied_password_stdin_empty_is_none(monkeypatch):
import io
monkeypatch.delenv(tp.SUPPLIED_PASSWORD_ENV, raising = False)
monkeypatch.setattr(sys, "stdin", io.StringIO(""))
assert tp.resolve_supplied_password("-") is None
def test_resolve_supplied_password_env(monkeypatch):
monkeypatch.setenv(tp.SUPPLIED_PASSWORD_ENV, "env-secret-pw")
assert tp.resolve_supplied_password("") == "env-secret-pw"
assert tp.resolve_supplied_password(None) == "env-secret-pw"
def test_resolve_supplied_password_literal_beats_env(monkeypatch):
import io
monkeypatch.setenv(tp.SUPPLIED_PASSWORD_ENV, "env-secret-pw")
assert tp.resolve_supplied_password("cli-wins-pw", out = io.StringIO()) == "cli-wins-pw"
def test_resolve_supplied_password_stdin_beats_env(monkeypatch):
# `--password -` reads stdin and short-circuits, so a set env var does not win.
import io
monkeypatch.setenv(tp.SUPPLIED_PASSWORD_ENV, "env-secret-pw")
monkeypatch.setattr(sys, "stdin", io.StringIO("stdin-wins-pw\n"))
assert tp.resolve_supplied_password("-") == "stdin-wins-pw"
def test_resolve_supplied_password_off_by_default(monkeypatch):
monkeypatch.delenv(tp.SUPPLIED_PASSWORD_ENV, raising = False)
assert tp.resolve_supplied_password("") is None
assert tp.resolve_supplied_password(None) is None

View file

@ -0,0 +1,405 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Pre-tunnel terminal password gate: never publish a public Cloudflare URL
while the seeded default admin password is active. Imports run.py directly,
so run under the Studio venv."""
from __future__ import annotations
import io
import re
import sys
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import run # noqa: E402
from auth import storage as auth_storage # noqa: E402
from auth import terminal_prompt # noqa: E402
from auth.terminal_prompt import should_prompt_password_change # noqa: E402
_GATE_KWARGS = dict(
host = "127.0.0.1",
secure = True,
api_only = False,
frontend_served = True,
)
# ── pure decision matrix ─────────────────────────────────────────────
@pytest.mark.parametrize(
"tunnel_will_start,requires_change,stdin_isatty,stderr_isatty,expected",
[
(True, True, True, True, True),
# Any missing precondition suppresses the prompt.
(False, True, True, True, False),
(True, False, True, True, False),
(True, True, False, True, False),
(True, True, True, False, False),
(False, False, False, False, False),
],
)
def test_should_prompt_matrix(
tunnel_will_start, requires_change, stdin_isatty, stderr_isatty, expected
):
assert (
should_prompt_password_change(
tunnel_will_start = tunnel_will_start,
requires_change = requires_change,
stdin_isatty = stdin_isatty,
stderr_isatty = stderr_isatty,
)
is expected
)
# ── _terminal_password_gate unit tests ───────────────────────────────
class _Stream(io.StringIO):
def __init__(self, isatty: bool):
super().__init__()
self._isatty = isatty
def isatty(self) -> bool:
return self._isatty
class _BrokenStream(io.StringIO):
"""Service-wrapper stand-in whose isatty() raises (closed stdin)."""
def isatty(self) -> bool:
raise ValueError("I/O operation on closed file")
def _patch_streams(monkeypatch, *, tty: bool) -> _Stream:
stderr = _Stream(isatty = tty)
monkeypatch.setattr(sys, "stdin", _Stream(isatty = tty))
monkeypatch.setattr(sys, "stderr", stderr)
return stderr
def _patch_seeded_admin(monkeypatch, *, requires_change: bool) -> None:
# The gate seeds the admin row itself (it can run before lifespan startup);
# tests fake both the seeding no-op and the flag.
monkeypatch.setattr(auth_storage, "ensure_default_admin", lambda: False)
monkeypatch.setattr(auth_storage, "requires_password_change", lambda u: requires_change)
def test_gate_skips_when_tunnel_off(monkeypatch):
# Short-circuits before touching auth storage at all.
def _boom(*a, **k):
raise AssertionError("storage must not be consulted when the tunnel is off")
monkeypatch.setattr(auth_storage, "requires_password_change", _boom)
monkeypatch.setattr(auth_storage, "ensure_default_admin", _boom)
assert run._terminal_password_gate(tunnel_will_start = False, **_GATE_KWARGS) == (True, False)
def test_gate_skips_when_password_already_changed(monkeypatch):
_patch_streams(monkeypatch, tty = True)
_patch_seeded_admin(monkeypatch, requires_change = False)
monkeypatch.setattr(
terminal_prompt,
"prompt_for_password_change",
lambda **k: pytest.fail("prompt must not run when no change is required"),
)
assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, False)
def test_gate_warns_and_proceeds_without_tty_when_deadline_arms(monkeypatch):
stderr = _patch_streams(monkeypatch, tty = False)
_patch_seeded_admin(monkeypatch, requires_change = True)
monkeypatch.delenv("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", raising = False)
monkeypatch.setattr(
terminal_prompt,
"prompt_for_password_change",
lambda **k: pytest.fail("prompt must not run without a tty"),
)
# Proceeds, but the public HTML must not auto-fill the default credential.
assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, True)
out = stderr.getvalue()
assert "default admin password is still active" in out
assert "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT" in out
# The seeded file may already be gone (the CLI parent deletes it before
# re-exec), so the warning must point at the reset-password recovery path
# instead of promising a file to read.
assert "reset-password" in out
assert ".bootstrap_password" not in out
def test_gate_fails_closed_without_tty_when_deadline_cannot_arm(monkeypatch):
# api-only launches never arm the bootstrap deadline, so a headless public
# launch with the default password has NO safeguard: refuse to start.
stderr = _patch_streams(monkeypatch, tty = False)
_patch_seeded_admin(monkeypatch, requires_change = True)
monkeypatch.delenv("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", raising = False)
kwargs = dict(_GATE_KWARGS)
kwargs["api_only"] = True
kwargs["frontend_served"] = False
assert run._terminal_password_gate(tunnel_will_start = True, **kwargs) == (False, False)
assert "Refusing to publish" in stderr.getvalue()
def test_gate_fails_closed_without_tty_when_deadline_disabled(monkeypatch):
stderr = _patch_streams(monkeypatch, tty = False)
_patch_seeded_admin(monkeypatch, requires_change = True)
monkeypatch.setenv("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", "0")
assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (False, False)
assert "Refusing to publish" in stderr.getvalue()
def test_gate_treats_broken_streams_as_non_interactive(monkeypatch):
# A closed/None stdin must take the headless path, not blow up.
stderr = _Stream(isatty = False)
monkeypatch.setattr(sys, "stdin", _BrokenStream())
monkeypatch.setattr(sys, "stderr", stderr)
_patch_seeded_admin(monkeypatch, requires_change = True)
monkeypatch.delenv("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", raising = False)
assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, True)
def test_gate_refusal_fails_closed(monkeypatch):
_patch_streams(monkeypatch, tty = True)
_patch_seeded_admin(monkeypatch, requires_change = True)
monkeypatch.setattr(terminal_prompt, "prompt_for_password_change", lambda **k: False)
assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (False, False)
def test_gate_success_applies_route_equivalent_change(monkeypatch):
_patch_streams(monkeypatch, tty = True)
calls = []
_patch_seeded_admin(monkeypatch, requires_change = True)
monkeypatch.setattr(
auth_storage,
"get_user_and_secret",
lambda u: ("salt", "hash", "jwt", True),
)
monkeypatch.setattr(
auth_storage,
"update_password",
lambda u, p, **kw: calls.append(("update", u, p, kw)),
)
def _fake_prompt(*, min_length, is_current_password, apply_change, out):
# The gate wires the policy constant and route-equivalent apply hook.
assert min_length == auth_storage.MIN_PASSWORD_LENGTH
# Wired to the real hash comparison: a wrong guess is rejected.
assert is_current_password("wrong-guess") is False
apply_change("brand-new-password")
return True
monkeypatch.setattr(terminal_prompt, "prompt_for_password_change", _fake_prompt)
assert run._terminal_password_gate(tunnel_will_start = True, **_GATE_KWARGS) == (True, True)
admin = auth_storage.DEFAULT_ADMIN_USERNAME
# One atomic call: refresh tokens revoked in the same transaction as the
# password commit (a separable follow-up delete can fail and leave a
# pre-change refresh token able to mint access tokens).
assert calls == [("update", admin, "brand-new-password", {"revoke_refresh_tokens": True})]
# ── ordering inside run_server (source-level, repo convention) ───────
def test_gate_runs_before_server_bind_in_source():
# The gate must run before the uvicorn socket binds: on a wildcard bind
# the served HTML injects the bootstrap credential for first login, so a
# pre-gate listener would hand out the default password mid-prompt.
src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
gate_call = src.index("_pw_proceed, _pw_drop_bootstrap = _terminal_password_gate(")
thread_start = src.index("thread.start()")
tunnel_start = src.index("_cloudflare_url = start_studio_tunnel(port)")
assert gate_call < thread_start < tunnel_start
# The fail-closed branch exits before any server exists.
refusal = src[gate_call:thread_start]
assert "sys.exit(1)" in refusal
def test_min_password_length_single_source():
# models/auth.py must reference the storage constant, not a literal.
models_src = (_BACKEND / "models" / "auth.py").read_text(encoding = "utf-8")
assert "MIN_PASSWORD_LENGTH" in models_src
assert not re.search(r"min_length\s*=\s*8\b", models_src)
assert auth_storage.MIN_PASSWORD_LENGTH == 8
def test_lifespan_honors_bootstrap_suppression_in_source():
# The lifespan runs AFTER the gate and re-reads the bootstrap password
# into app.state; without the suppress flag it would overwrite the gate's
# None and the public HTML would inject the default credential again.
main_src = (_BACKEND / "main.py").read_text(encoding = "utf-8")
assert "suppress_bootstrap_injection" in main_src
# Every lifespan capture of the bootstrap password must be flag-guarded.
for line in main_src.splitlines():
if "storage.get_bootstrap_password()" in line and "=" in line:
assert "_suppress_bootstrap" in line, line
run_src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
assert "app.state.suppress_bootstrap_injection = True" in run_src
def test_clear_bootstrap_password_truncates_when_unlink_fails(monkeypatch, tmp_path):
# If the file cannot be unlinked (Windows AV / read-only auth dir), clear must
# truncate it so its stale plaintext cannot be re-seeded by
# generate_bootstrap_password() after a later reset-password deletes auth.db,
# which would re-validate the revoked bootstrap password.
import pathlib
pw_path = tmp_path / ".bootstrap_password"
pw_path.write_text("old-diceware-passphrase")
monkeypatch.setattr(auth_storage, "_BOOTSTRAP_PW_PATH", pw_path)
monkeypatch.setattr(auth_storage, "_bootstrap_password", "old-diceware-passphrase")
_real_unlink = pathlib.Path.unlink
def _boom(self, *a, **k):
if self == pw_path:
raise OSError("locked")
return _real_unlink(self, *a, **k)
monkeypatch.setattr(pathlib.Path, "unlink", _boom)
auth_storage.clear_bootstrap_password()
assert pw_path.exists() # unlink failed
assert pw_path.read_text() == "" # but truncated -> no reusable plaintext
# The stale value must not load back (empty file -> None), so a later re-seed
# generates fresh rather than resurrecting the revoked credential.
monkeypatch.setattr(auth_storage, "_bootstrap_password", None)
assert auth_storage._load_bootstrap_password() is None
def test_clear_bootstrap_password_warns_truthfully_when_not_cleared(monkeypatch, tmp_path, capsys):
# If the file can be neither unlinked NOR truncated, the stale plaintext stays
# on disk. The warning must NOT claim it was made unreusable (Codex 3571888584):
# it must say it could not be cleared and ask the user to remove it manually.
import pathlib
pw_path = tmp_path / ".bootstrap_password"
pw_path.write_text("old-diceware-passphrase")
monkeypatch.setattr(auth_storage, "_BOOTSTRAP_PW_PATH", pw_path)
monkeypatch.setattr(auth_storage, "_bootstrap_password", "old-diceware-passphrase")
_real_unlink = pathlib.Path.unlink
_real_write_text = pathlib.Path.write_text
def _boom_unlink(self, *a, **k):
if self == pw_path:
raise OSError("locked")
return _real_unlink(self, *a, **k)
def _boom_write_text(self, *a, **k):
if self == pw_path:
raise OSError("read-only")
return _real_write_text(self, *a, **k)
monkeypatch.setattr(pathlib.Path, "unlink", _boom_unlink)
monkeypatch.setattr(pathlib.Path, "write_text", _boom_write_text)
auth_storage.clear_bootstrap_password()
# The stale plaintext survives untouched.
assert pw_path.read_text() == "old-diceware-passphrase"
warning = capsys.readouterr().err.lower()
assert "could not delete or clear" in warning
assert "still on disk" in warning
assert "remove it manually" in warning
# Must not falsely claim the contents were cleared (the bug being fixed).
assert "cleared its contents" not in warning
# ── _apply_supplied_password: non-interactive initial password (direct run.py) ──
def _seed_stub_admin(
monkeypatch,
*,
requires_change,
bootstrap_pw = "bootstrap-secret",
):
"""Stub storage so _apply_supplied_password sees a seeded admin whose current
password is ``bootstrap_pw`` and whose must-change flag is ``requires_change``;
return the recorded update_password calls."""
from auth import hashing
salt, pwd_hash = hashing.hash_password(bootstrap_pw)
monkeypatch.setattr(auth_storage, "ensure_default_admin", lambda: False)
monkeypatch.setattr(auth_storage, "requires_password_change", lambda u: requires_change)
monkeypatch.setattr(
auth_storage, "get_user_and_secret", lambda u: (salt, pwd_hash, "jwt", requires_change)
)
calls = []
monkeypatch.setattr(
auth_storage, "update_password", lambda u, p, **kw: calls.append((u, p, kw))
)
return calls
def test_apply_supplied_password_sets_initial(monkeypatch):
calls = _seed_stub_admin(monkeypatch, requires_change = True)
monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "brand-new-password")
run._apply_supplied_password(None) # resolves from the env var
admin = auth_storage.DEFAULT_ADMIN_USERNAME
assert calls == [(admin, "brand-new-password", {"revoke_refresh_tokens": True})]
def test_apply_supplied_password_off_is_noop(monkeypatch):
calls = _seed_stub_admin(monkeypatch, requires_change = True)
monkeypatch.delenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, raising = False)
run._apply_supplied_password(None)
run._apply_supplied_password("")
assert calls == []
def test_apply_supplied_password_already_set_fails_closed(monkeypatch):
calls = _seed_stub_admin(monkeypatch, requires_change = False)
monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "brand-new-password")
with pytest.raises(SystemExit) as exc:
run._apply_supplied_password(None)
assert exc.value.code == 1
assert calls == [] # never overrides an existing password
def test_apply_supplied_password_too_short_fails_closed(monkeypatch):
calls = _seed_stub_admin(monkeypatch, requires_change = True)
monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "short")
with pytest.raises(SystemExit) as exc:
run._apply_supplied_password(None)
assert exc.value.code == 1
assert calls == []
def test_apply_supplied_password_must_differ_fails_closed(monkeypatch):
calls = _seed_stub_admin(monkeypatch, requires_change = True, bootstrap_pw = "bootstrap-secret")
monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "bootstrap-secret")
with pytest.raises(SystemExit) as exc:
run._apply_supplied_password(None)
assert exc.value.code == 1
assert calls == []
def test_apply_supplied_password_strips_env_from_subprocess_environment(monkeypatch):
# The plaintext password must not linger in os.environ: run_server later spawns
# cloudflared/llama-server/code-exec tools that would otherwise inherit it (also
# readable via /proc/PID/environ). The direct-run.py path pops it itself; the CLI
# pops it before re-exec. Assert the pop happens on the apply path...
_seed_stub_admin(monkeypatch, requires_change = True)
monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "brand-new-password")
run._apply_supplied_password(None)
assert terminal_prompt.SUPPLIED_PASSWORD_ENV not in run.os.environ
def test_apply_supplied_password_strips_env_even_when_literal_wins(monkeypatch):
# A literal --password wins over the env var, but a stale env value would still
# leak to subprocesses; the unconditional pop must clear it regardless of source.
_seed_stub_admin(monkeypatch, requires_change = True)
monkeypatch.setenv(terminal_prompt.SUPPLIED_PASSWORD_ENV, "env-should-be-stripped")
run._apply_supplied_password("literal-new-password")
assert terminal_prompt.SUPPLIED_PASSWORD_ENV not in run.os.environ

File diff suppressed because it is too large Load diff

View file

@ -2592,6 +2592,50 @@ class TestLoopBasic:
assert tool_starts[0]["arguments"] == {}
assert "<!doctype html>" in tool_starts[1]["arguments"]["code"]
def test_render_html_auto_mode_static_runs_without_prompt(self):
"""permission_mode="auto" ships confirm_tool_calls=true. render_html is no
longer unconditionally safe (a networked canvas must ask), so its early
provisional card is suppressed under the confirm gate; a static canvas is
still classified safe and runs without an approval prompt."""
exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
turn_iter = iter(
[
[
"<function=render_html>",
"<parameter=code><!doctype html><html>",
"<body>Hi</body></html></parameter></function>",
],
["Done."],
]
)
def _gen(_messages):
chunks = next(turn_iter)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
loop = run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "make html"}],
tools = [{"type": "function", "function": {"name": "render_html"}}],
execute_tool = exec_fn,
confirm_tool_calls = True,
permission_mode = "auto",
session_id = "sess",
max_tool_iterations = 3,
)
events = _collect_events(loop)
tool_starts = [e for e in events if e["type"] == "tool_start"]
# No early provisional card under the auto confirm gate; just the real call.
assert len(tool_starts) == 1
assert tool_starts[0]["tool_name"] == "render_html"
assert "<!doctype html>" in tool_starts[0]["arguments"]["code"]
# A static canvas is classified safe, so it runs without an approval gate.
assert tool_starts[0].get("awaiting_confirmation") in (False, None)
def test_render_html_provisional_card_closed_on_generator_exception(self):
"""If the model generator raises mid-stream after a provisional render_html
card was surfaced, the loop must close that card as errored before the
@ -3674,6 +3718,26 @@ class TestGuardrails:
assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events)
assert exec_fn.calls == []
def test_auto_mode_still_runs_rag_autoinject(self, monkeypatch):
# "auto" sends confirm_tool_calls=true so unsafe calls gate, but the
# safe search_knowledge_base retrieval never gates, so autoinject must
# still run (unlike ask mode above).
ran = {"called": False}
def fake_autoinject(*_args, **_kwargs):
ran["called"] = True
return None
monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fake_autoinject)
loop, _exec_fn = _make_loop(
turns = [["plain answer"]],
confirm_tool_calls = True,
permission_mode = "auto",
rag_scope = {"thread_id": "t1"},
)
_collect_events(loop)
assert ran["called"] is True
def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self):
turns = iter(
[

View file

@ -21,7 +21,7 @@ from run import _cloudflare_tunnel_should_start as should_start # noqa: E402
@pytest.mark.parametrize(
"cloudflare,host,secure,api_only,is_colab,expected",
[
# Non-secure wildcard binds tunnel by default.
# Non-secure wildcard binds tunnel only when --cloudflare is passed (True).
(True, "0.0.0.0", False, False, False, True),
(True, "::", False, False, False, True),
(True, "127.0.0.1", False, False, False, False),
@ -33,6 +33,10 @@ from run import _cloudflare_tunnel_should_start as should_start # noqa: E402
(False, "0.0.0.0", False, False, False, False),
(False, "::", False, False, False, False),
(False, "127.0.0.1", True, False, False, False),
# Unset (None, no flag) behaves as off for non-secure binds.
(None, "0.0.0.0", False, False, False, False),
(None, "::", False, False, False, False),
(None, "127.0.0.1", False, False, False, False),
# Non-secure api-only never tunnels (Tauri).
(True, "0.0.0.0", False, True, False, False),
(True, "::", False, True, False, False),
@ -155,11 +159,12 @@ def test_startup_output_emits_disabled_notice(capsys, monkeypatch):
def test_run_server_rejects_secure_without_cloudflare():
# Direct backend callers (not just the CLI) must reject the contradictory combo.
# Direct backend callers (not just the CLI) must reject the contradictory
# combo: --secure asks for the tunnel, --no-cloudflare (cloudflare=False) forbids it.
import run
with pytest.raises(SystemExit) as exc:
run.run_server(secure = True, cloudflare = False)
assert "A secure Cloudflare link is not allowed" in str(exc.value)
assert "do not combine it with --no-cloudflare" in str(exc.value)
def test_failclosed_message_present_in_source():

View file

@ -0,0 +1,145 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""_shutdown_subprocess returns whether the worker actually died, and preserves the
live handle when it survives terminate/kill.
A GPU worker wedged in an uninterruptible CUDA syscall can outlive SIGKILL. If shutdown
nulled its handle anyway, is_worker_alive() would report False and the pre-swap liveness
guard would let the destructive .venv_t5_latest rename proceed while a live worker still
holds sidecar transformers modules (breaking the rename on Windows). The methods must keep
the handle and return False so callers can refuse the swap.
"""
import pytest
from core.export.orchestrator import ExportOrchestrator
from core.inference.orchestrator import InferenceOrchestrator
class _FakeProc:
"""A subprocess handle that dies only on the requested step (or never)."""
def __init__(self, dies_on = None):
self._alive = True
self._dies_on = dies_on # None | "join" | "terminate" | "kill"
def is_alive(self):
return self._alive
def join(self, timeout = None):
if self._dies_on == "join":
self._alive = False
def terminate(self):
if self._dies_on == "terminate":
self._alive = False
def kill(self):
if self._dies_on == "kill":
self._alive = False
def _bare_inference():
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
o._stop_dispatcher = lambda: None
o._cancel_generation = lambda: None
o._drain_queue = lambda: []
class _Q:
def put(self, *a, **k):
pass
o._cmd_queue = _Q()
o._resp_queue = _Q()
o._cancel_event = None
o._drain_event = None
return o
def _bare_export():
o = ExportOrchestrator.__new__(ExportOrchestrator)
o._drain_queue = lambda: []
class _Q:
def put(self, *a, **k):
pass
o._cmd_queue = _Q()
o._resp_queue = _Q()
return o
@pytest.fixture(autouse = True)
def _no_sleep(monkeypatch):
# _shutdown_subprocess sleeps 0.5s after cancelling; keep the tests instant.
import core.inference.orchestrator as inf_mod
monkeypatch.setattr(inf_mod.time, "sleep", lambda *_a, **_k: None)
class TestInferenceShutdownReturn:
def test_worker_that_dies_returns_true_and_clears_handle(self):
o = _bare_inference()
o._proc = _FakeProc(dies_on = "terminate")
assert o._shutdown_subprocess(timeout = 0.01) is True
assert o._proc is None
assert o.is_worker_alive() is False
def test_survivor_returns_false_and_keeps_handle(self):
o = _bare_inference()
o._proc = _FakeProc(dies_on = None) # outlives terminate AND kill
assert o._shutdown_subprocess(timeout = 0.01) is False
assert o._proc is not None
# is_worker_alive stays truthful, so the pre-swap guard can refuse the swap.
assert o.is_worker_alive() is True
def test_already_dead_returns_true(self):
o = _bare_inference()
o._proc = _FakeProc(dies_on = "join")
o._proc._alive = False
assert o._shutdown_subprocess(timeout = 0.01) is True
assert o._proc is None
class TestExportShutdownReturn:
def test_worker_that_dies_returns_true_and_clears_handle(self):
o = _bare_export()
o._proc = _FakeProc(dies_on = "terminate")
assert o._shutdown_subprocess(timeout = 0.01) is True
assert o._proc is None
assert o.is_worker_alive() is False
def test_survivor_returns_false_and_keeps_handle(self):
o = _bare_export()
o._proc = _FakeProc(dies_on = None)
assert o._shutdown_subprocess(timeout = 0.01) is False
assert o._proc is not None
assert o.is_worker_alive() is True
class TestSpawnPathsHonorFailedShutdown:
"""A fresh-load path must not spawn a second worker over one that outlived
terminate/kill: the survivor still holds GPU memory and its handle would be lost."""
def test_export_load_checkpoint_aborts_when_worker_survives(self, monkeypatch):
import threading
import utils.transformers_version as tv
o = ExportOrchestrator.__new__(ExportOrchestrator)
o._lock = threading.RLock()
o._proc = _FakeProc(dies_on = None) # survivor
o.clear_logs = lambda: None
o._cancel_requested = False
o._active_op_kind = None
o._export_active = False
o._ensure_subprocess_alive = lambda: True
o._shutdown_subprocess = lambda *a, **k: False
o._spawn_subprocess = lambda cfg: pytest.fail("must not spawn over a live survivor")
o._record_op_finished = lambda *a, **k: None
monkeypatch.setattr(tv, "sidecar_swap_in_progress", lambda: False)
ok, msg = o.load_checkpoint(checkpoint_path = "ckpt")
assert ok is False
assert "did not exit" in msg
# The finally cleared the op flags even though we returned early.
assert o._export_active is False

View file

@ -0,0 +1,824 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Stop-watchdog escalation for a stuck training stop.
A save-stop signals the worker and waits for it to save and exit. On some platforms the
worker saves but then wedges in post-save GPU/driver teardown and never exits, leaving the
run stuck in "Stopping..." forever. These tests pin the bounded recovery: the watchdog
escalates to force_terminate() a short grace after "complete" (save done) or after an
absolute timeout (hang during save), and never force-kills a worker that exits cleanly.
Fakes only; no GPU, network, or subprocess.
"""
from __future__ import annotations
import contextlib
import logging
import queue
import sys
import threading
import time
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub the heavy module-level imports of core/training/training.py so it imports
# under CPU-only/no-network, then restore them (see the restore loop below).
_SAVED: dict = {}
def _stub(name, mod):
_SAVED[name] = sys.modules.get(name)
sys.modules[name] = mod
_lg = _types.ModuleType("loggers")
_lg.get_logger = lambda name: logging.getLogger(name)
_stub("loggers", _lg)
_stub("structlog", _types.ModuleType("structlog"))
_mpl = _types.ModuleType("matplotlib")
_plt = _types.ModuleType("matplotlib.pyplot")
_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation
_mpl.pyplot = _plt
_stub("matplotlib", _mpl)
_stub("matplotlib.pyplot", _plt)
_hw = _types.ModuleType("utils.hardware")
_hw.prepare_gpu_selection = lambda *a, **k: (None, None)
_stub("utils.hardware", _hw)
_npl = _types.ModuleType("utils.native_path_leases")
_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext()
_npl.run_without_native_path_secret = lambda fn: fn
_stub("utils.native_path_leases", _npl)
_pth = _types.ModuleType("utils.paths")
_pth.outputs_root = lambda *a, **k: "/tmp/outputs"
_stub("utils.paths", _pth)
# Whether core.training.training was already imported before this file ran; only
# evict it below if we were the one to create the (stub-bound) module instance.
_TRAINING_PRE_IMPORTED = "core.training.training" in sys.modules
from core.training.training import TrainingBackend
# Restore every stubbed module so this file never pollutes the shared session.
for _name in (
"loggers",
"structlog",
"matplotlib",
"matplotlib.pyplot",
"utils.hardware",
"utils.native_path_leases",
"utils.paths",
):
_prev = _SAVED.get(_name)
if _prev is None:
sys.modules.pop(_name, None)
else:
sys.modules[_name] = _prev
if not _TRAINING_PRE_IMPORTED:
sys.modules.pop("core.training.training", None)
sys.modules.pop("core.training", None)
# The module globals hold the escalation timeouts and are the watchdog's own
# namespace; patch them here so tests run in well under a second.
_G = TrainingBackend._stop_watchdog_loop.__globals__
class _FakeProc:
"""A subprocess handle whose liveness and kill calls the test observes."""
def __init__(self, alive: bool = True):
self._alive = alive
self.pid = 4321
self.terminated = False
self.killed = False
def is_alive(self):
return self._alive
def terminate(self):
self.terminated = True
def kill(self):
self.killed = True
def join(self, timeout = None):
pass
def _wait_until(predicate, timeout = 5.0):
deadline = time.time() + timeout
while time.time() < deadline:
if predicate():
return True
time.sleep(0.01)
return predicate()
def _record_force_terminate(monkeypatch, b):
"""Replace force_terminate + escalation finalize with recorders (no DB/OS)."""
calls: list = []
monkeypatch.setattr(b, "force_terminate", lambda target_proc = None: calls.append("force"))
monkeypatch.setattr(
b,
"_finalize_stopped_after_escalation",
lambda target_proc = None, watched_job_id = None: calls.append("final"),
)
return calls
# ----------------------------------------------------------------------------
# (a) Escalate a short grace after "complete" (save done) if still alive.
# ----------------------------------------------------------------------------
def test_watchdog_escalates_after_grace_once_complete_seen(monkeypatch):
monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # ensure grace, not timeout, fires
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
proc = _FakeProc(alive = True)
b._proc = proc
b._complete_seen.set() # worker reported "complete" -> save is done
b._start_stop_watchdog(cancel = False)
assert _wait_until(
lambda: calls == ["force", "final"]
), "watchdog must force_terminate a worker still alive after the post-save grace"
b._stop_watchdog.join(timeout = 5)
# ----------------------------------------------------------------------------
# (b) The absolute cap is a last-resort backstop, not a save killer.
# ----------------------------------------------------------------------------
def test_watchdog_does_not_kill_save_still_saving_within_window(monkeypatch):
# save=True, no "complete" yet: a slow save in progress must not be force-killed
# inside the (long) absolute window.
monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0)
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
proc = _FakeProc(alive = True)
b._proc = proc
b._start_stop_watchdog(cancel = False)
time.sleep(0.3)
assert calls == [], "an in-progress save must not be killed within the absolute window"
assert b._stop_watchdog.is_alive()
proc._alive = False
b._stop_watchdog.join(timeout = 5)
def test_watchdog_backstop_fires_for_save_after_absolute_timeout(monkeypatch):
# Past the long save=True cap with no completion: force-terminate as last resort.
monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) # never trips (no complete)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 0.05)
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
b._proc = _FakeProc(alive = True)
b._start_stop_watchdog(cancel = False)
assert _wait_until(
lambda: calls == ["force", "final"]
), "the absolute backstop must force_terminate a save that never completes"
b._stop_watchdog.join(timeout = 5)
def test_cancel_uses_shorter_absolute_timeout(monkeypatch):
# A cancel has nothing to save, so it escalates on the shorter cancel cap even before
# the long save cap elapses.
monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # save cap would not fire
monkeypatch.setitem(_G, "_CANCEL_TIMEOUT_S", 0.05)
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
b._proc = _FakeProc(alive = True)
b._start_stop_watchdog(cancel = True)
assert _wait_until(
lambda: calls == ["force", "final"]
), "a cancel must escalate on the shorter cancel timeout"
b._stop_watchdog.join(timeout = 5)
# ----------------------------------------------------------------------------
# (c) No force-kill when the worker exits cleanly and promptly.
# ----------------------------------------------------------------------------
def test_watchdog_no_op_on_clean_quick_exit(monkeypatch):
monkeypatch.setitem(_G, "_STOP_GRACE_S", 5.0)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 10.0)
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
proc = _FakeProc(alive = True)
b._proc = proc
b._complete_seen.set() # save done; worker is about to exit on its own
b._start_stop_watchdog(cancel = False)
# Worker exits promptly, well before the grace period elapses.
time.sleep(0.1)
proc._alive = False
b._stop_watchdog.join(timeout = 5)
assert not b._stop_watchdog.is_alive()
assert calls == [], "a clean quick exit must not trigger force_terminate"
def test_watchdog_no_op_when_worker_superseded(monkeypatch):
# A stale watchdog from a prior run must never kill a new run's worker: once
# self._proc is replaced, it exits silently.
monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 0.05)
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
old_proc = _FakeProc(alive = True)
b._proc = old_proc
b._complete_seen.set()
b._start_stop_watchdog(cancel = False)
# A new run takes over the handle before the grace elapses.
b._proc = _FakeProc(alive = True)
b._stop_watchdog.join(timeout = 5)
assert calls == [], "watchdog must not force_terminate a superseded worker"
def test_new_run_gets_its_own_watchdog(monkeypatch):
# A stale watchdog sleeping on an old proc must not stop a new run's stop from
# creating its own watcher.
monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0)
b = TrainingBackend()
_record_force_terminate(monkeypatch, b)
old_proc = _FakeProc(alive = True)
b._proc = old_proc
b._start_stop_watchdog(cancel = False)
first_wd = b._stop_watchdog
# New run: fresh worker replaces the handle; its stop must get a new watcher
# even though the old (superseded) watchdog is still alive.
new_proc = _FakeProc(alive = True)
b._proc = new_proc
b._start_stop_watchdog(cancel = False)
second_wd = b._stop_watchdog
try:
assert first_wd.is_alive()
assert second_wd is not first_wd, "a new run must get its own watchdog"
assert b._stop_watchdog_proc is new_proc
finally:
old_proc._alive = False
new_proc._alive = False
first_wd.join(timeout = 5)
second_wd.join(timeout = 5)
def test_force_terminate_targets_only_captured_proc():
# Superseded: force_terminate(target) must not touch a different current worker.
b = TrainingBackend()
old_proc = _FakeProc(alive = True)
new_proc = _FakeProc(alive = True)
b._proc = new_proc
b.force_terminate(target_proc = old_proc)
assert new_proc.terminated is False, "must not terminate the new run's worker"
assert old_proc.terminated is False, "must not terminate a handle that is not current"
# Matching: the captured handle is the current worker, so it is terminated.
p = _FakeProc(alive = True)
b._proc = p
b.force_terminate(target_proc = p)
assert p.terminated is True
# ----------------------------------------------------------------------------
# Post-escalation finalize leaves the parent ready for a new run.
# ----------------------------------------------------------------------------
def test_finalize_runs_even_if_force_terminate_raises(monkeypatch):
# A wedged child can make force_terminate() raise; finalize must still run so the
# run does not stay stuck in "Stopping...".
monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0)
b = TrainingBackend()
def _boom(target_proc = None):
raise RuntimeError("kill() failed on wedged child")
finalized: list = []
monkeypatch.setattr(b, "force_terminate", _boom)
monkeypatch.setattr(
b,
"_finalize_stopped_after_escalation",
lambda target_proc = None, watched_job_id = None: finalized.append(True),
)
b._proc = _FakeProc(alive = True)
b._complete_seen.set()
b._start_stop_watchdog(cancel = False)
assert _wait_until(
lambda: finalized == [True]
), "finalize must run even when force_terminate raises"
b._stop_watchdog.join(timeout = 5)
def test_finalize_after_escalation_clears_state(monkeypatch):
# Even if the OS never reaps the wedged worker, the parent must report the run
# stopped so the UI leaves "Stopping..." and a new run can start.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
b._proc = _FakeProc(alive = True) # wedged: still reports alive
b._should_stop = True
b.current_job_id = "job_c"
b._db_run_created = True
b._progress.is_training = True
b._finalize_stopped_after_escalation(watched_job_id = "job_c")
assert b._proc is None, "the wedged handle must be dropped so is_training_active clears"
assert b._progress.is_training is False
assert b._progress.status_message == "Training stopped."
assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id"
assert b.is_training_active() is False
def test_finalize_after_escalation_preserves_output_dir(monkeypatch):
# A save-stop that already emitted "complete" has the checkpoint dir; run history
# must record it even if the watchdog wins the finalize race against the pump.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
b._proc = _FakeProc(alive = True)
b._should_stop = True
b.current_job_id = "job_c"
b._db_run_created = True
b._output_dir = "/tmp/outputs/run-123"
b._finalize_stopped_after_escalation(watched_job_id = "job_c")
# _finish_stopped_run(run_id, output_dir, batch, final_step, final_loss, duration, loss_history)
assert finstop and finstop[0][0] == "job_c"
assert finstop[0][1] == "/tmp/outputs/run-123"
def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch):
# No worker -> nothing to escalate; the watchdog must not spawn.
b = TrainingBackend()
b._proc = None
assert b.stop_training(save = True) is True
assert b._stop_watchdog is None
# ----------------------------------------------------------------------------
# (d) A stale watchdog must never clobber a run that replaced its worker.
# ----------------------------------------------------------------------------
def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch):
# A /start can slip in while the watchdog force-terminates the old worker
# (is_training_active() is False once _should_stop is set and the old proc is dead).
# The escalation finalize must then leave the NEW run untouched, not drop its handle.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
old_proc = _FakeProc(alive = False) # force-terminated worker we were watching
new_proc = _FakeProc(alive = True) # a new run already took over
b._proc = new_proc
b.current_job_id = "job_new"
b._db_run_created = True
b._progress.is_training = True
b._finalize_stopped_after_escalation(target_proc = old_proc)
assert b._proc is new_proc, "must not drop the new run's handle"
assert b._progress.is_training is True, "must not mark the new run stopped"
assert finstop == [], "must not finalize the new run in the DB"
def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch):
# Common case: the watched worker is still current, so finalize proceeds and
# finalizes the captured run by id.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
proc = _FakeProc(alive = False)
b._proc = proc
b.current_job_id = "job_a"
b._db_run_created = True
b._progress.is_training = True
b._finalize_stopped_after_escalation(target_proc = proc, watched_job_id = "job_a")
assert b._proc is None
assert b._progress.is_training is False
assert finstop and finstop[0][0] == "job_a", "must finalize the captured run by id"
def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypatch):
# start_training updates current_job_id BEFORE it installs the new _proc, so a stale
# watchdog can enter while _proc is still the old (dead) handle. The job-id guard must
# catch this even though the proc-only guard would not.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet
b._proc = old_proc # still the old handle (== target), so proc guard would pass
b.current_job_id = "job_new" # but the new run already claimed the job id
b._db_run_created = True
b._progress.is_training = True
b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old")
assert b._proc is old_proc, "must not drop the handle during a new run's startup"
assert b._progress.is_training is True, "must not mark the starting run stopped"
assert finstop == [], "must not finalize while a new run is starting up"
# ----------------------------------------------------------------------------
# (e) A later cancel (save=False) tightens an in-flight save watchdog.
# ----------------------------------------------------------------------------
def test_later_cancel_tightens_watchdog_timeout(monkeypatch):
monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) # never trips (no complete)
monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # save cap would not fire
monkeypatch.setitem(_G, "_CANCEL_TIMEOUT_S", 0.05)
b = TrainingBackend()
calls = _record_force_terminate(monkeypatch, b)
b._proc = _FakeProc(alive = True)
b._start_stop_watchdog(cancel = False) # started as a save-stop with the long cap
time.sleep(0.15)
assert calls == [], "a save-stop must not escalate on the short cancel cap yet"
# The user now cancels the in-flight stop: the watchdog must tighten its cap.
b._cancel_requested = True
assert _wait_until(
lambda: calls == ["force", "final"]
), "a later cancel must tighten the watchdog to the shorter cancel cap"
b._stop_watchdog.join(timeout = 5)
# ----------------------------------------------------------------------------
# (f) DB finalize/flush are safe when the watchdog and pump race (see Item 4).
# ----------------------------------------------------------------------------
def _install_fake_db(monkeypatch):
"""Stub storage.studio_db + utils.downsample so the real DB helpers run without
SQLite. Returns the recorder dict."""
recs = {"created": [], "finished": [], "inserted": [], "insert_ids": [], "progress_ids": []}
fake_storage = _types.ModuleType("storage")
fake_db = _types.ModuleType("storage.studio_db")
fake_db.create_run = lambda **kw: recs["created"].append(kw)
fake_db.finish_run = lambda **kw: recs["finished"].append(kw)
fake_db.insert_metrics_batch = lambda job_id, batch: (
recs["inserted"].extend(batch),
recs["insert_ids"].append(job_id),
)
fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id"))
fake_storage.studio_db = fake_db
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
fake_ds = _types.ModuleType("utils.downsample")
fake_ds.downsample = lambda seq, n: list(seq)[:n]
monkeypatch.setitem(sys.modules, "utils.downsample", fake_ds)
return recs
def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch):
# The watchdog and pump can both finalize; only one call may reach finish_run.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_x"
b._db_run_created = True
b._run_finalized = False
start = threading.Barrier(8)
def worker():
start.wait()
b._finalize_run_in_db(status = "stopped")
threads = [threading.Thread(target = worker) for _ in range(8)]
for t in threads:
t.start()
for t in threads:
t.join(timeout = 5)
assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}"
assert b._run_finalized is True
def test_finalize_run_in_db_no_ops_on_job_mismatch(monkeypatch):
# A finalize captured for an old job must not finalize the run that replaced it.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_new"
b._db_run_created = True
b._run_finalized = False
b._finalize_run_in_db(status = "stopped", expected_job_id = "job_old")
assert recs["finished"] == [], "a superseded job id must not finalize the current run"
assert b._run_finalized is False
def test_concurrent_flush_claims_each_metric_once(monkeypatch):
# Concurrent flushes (pump periodic flush vs watchdog finalize flush) must not
# double-remove or drop buffered metrics.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_y"
b._db_run_created = True
b._metric_buffer[:] = [{"step": i} for i in range(200)]
start = threading.Barrier(6)
def worker():
start.wait()
for _ in range(50):
b._flush_metrics_to_db()
threads = [threading.Thread(target = worker) for _ in range(6)]
for t in threads:
t.start()
for t in threads:
t.join(timeout = 5)
b._flush_metrics_to_db() # drain any remainder
steps = sorted(m["step"] for m in recs["inserted"])
assert steps == list(range(200)), "each metric must be inserted exactly once"
assert b._metric_buffer == [], "the buffer must be fully drained"
def test_flush_pins_to_passed_run_id(monkeypatch):
# A finalizer flushes to the run it captured, even if a new /start has already
# changed current_job_id.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_new" # a new run is already live
b._db_run_created = True
b._metric_buffer[:] = [{"step": 1}, {"step": 2}]
b._flush_metrics_to_db(run_id = "job_old")
assert recs["insert_ids"] == ["job_old"], "metrics must go to the captured run, not the new one"
assert recs["progress_ids"] == ["job_old"]
def test_finalize_uses_snapshot_run_id_across_new_run(monkeypatch):
# If a new /start changes current_job_id after the finalize claim but before the DB
# writes, finish_run must still target the run captured under the lock.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_x"
b._db_run_created = True
b._run_finalized = False
def hijack(run_id = None):
# Simulate a new run taking over during the flush (after the finalize claim).
b.current_job_id = "job_y"
monkeypatch.setattr(b, "_flush_metrics_to_db", hijack)
b._finalize_run_in_db(status = "stopped", expected_job_id = "job_x")
assert [f["id"] for f in recs["finished"]] == [
"job_x"
], "finish_run must target the captured run, not the run that replaced it"
# ----------------------------------------------------------------------------
# (g) DB row creation must not be published before the insert commits.
# ----------------------------------------------------------------------------
def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch):
# _db_run_created must stay False while create_run is in flight, so a concurrent
# finalize can't run finish_run (an UPDATE) against a not-yet-inserted row.
b = TrainingBackend()
b.current_job_id = "job_z"
b._db_config = {"model_name": "m"}
observed: dict = {}
fake_storage = _types.ModuleType("storage")
fake_db = _types.ModuleType("storage.studio_db")
def _create(**kw):
observed["flag_during_create"] = b._db_run_created
observed["in_progress_during_create"] = b._db_create_in_progress
fake_db.create_run = _create
fake_storage.studio_db = fake_db
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
b._ensure_db_run_created()
assert observed["flag_during_create"] is False, "flag must not be published before insert"
assert observed["in_progress_during_create"] is True
assert b._db_run_created is True, "flag must be published after a successful insert"
assert b._db_create_in_progress is False
def test_ensure_db_run_created_stays_unpublished_on_failure(monkeypatch):
# If create_run raises, neither flag stays set, so a later caller can retry.
b = TrainingBackend()
b.current_job_id = "job_z"
b._db_config = {"model_name": "m"}
fake_storage = _types.ModuleType("storage")
fake_db = _types.ModuleType("storage.studio_db")
def _boom_create(**kw):
raise RuntimeError("insert failed")
fake_db.create_run = _boom_create
fake_storage.studio_db = fake_db
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
b._ensure_db_run_created()
assert b._db_run_created is False, "a failed insert must not publish the row as created"
assert b._db_create_in_progress is False, "the in-progress flag must be cleared on failure"
def test_ensure_db_run_created_does_not_publish_for_a_new_run(monkeypatch):
# A killed worker lets a new /start proceed while the watchdog is still creating the old
# run's row. The stale create must not publish the backend-wide flags against the new
# current_job_id, or the new run would skip inserting its own row.
b = TrainingBackend()
b.current_job_id = "job_old"
b._db_config = {"model_name": "m"}
b._db_run_created = False
b._db_create_in_progress = False
fake_storage = _types.ModuleType("storage")
fake_db = _types.ModuleType("storage.studio_db")
def _create(**kw):
b.current_job_id = "job_new" # a new run takes over during the slow create
fake_db.create_run = _create
fake_storage.studio_db = fake_db
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
b._ensure_db_run_created()
assert b._db_run_created is False, "must not publish the created flag against the new run"
# The stale claim is left for start_training to reset, not satisfied for the new run.
assert b._db_create_in_progress is True, "must not clear the claim once the run is not current"
# ----------------------------------------------------------------------------
# (h) The escalation finalizes the watched run by id (so it is never left running).
# ----------------------------------------------------------------------------
def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch):
# Exercise the real _finish_stopped_run against a fake DB. The watched run is finalized
# by its captured id with its buffered metrics, so a new run that starts in the gap
# after the backend goes idle can never leave the stopped run recorded running.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_old"
b._db_run_created = True
b._proc = _FakeProc(alive = False)
b._progress.is_training = True
b._progress.step = 42
b._metric_buffer[:] = [{"step": 41}, {"step": 42}]
b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old")
assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id"
assert recs["finished"][0]["status"] == "stopped"
assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run"
assert b._metric_buffer == [], "the captured batch must be drained"
def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch):
# If the row does not exist and cannot be created here (no db_config, or the pump is
# mid-create), the escalation must not claim _run_finalized or call _finish_stopped_run,
# so the pump's create-then-finalize records the run. Parent state still clears.
b = TrainingBackend()
called: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a))
b._proc = _FakeProc(alive = False)
b.current_job_id = "job_q"
b._db_run_created = False # row not created yet
b._db_config = None # ... and cannot be created here
b._run_finalized = False
b._progress.is_training = True
b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_q")
assert called == [], "must not finalize when the row can't be established here"
assert b._run_finalized is False, "must not claim the finalize the pump still owes"
assert b._progress.is_training is False, "parent state must still clear so the UI unsticks"
assert b._proc is None
def test_escalation_creates_row_then_finalizes_when_start_create_failed(monkeypatch):
# A wedged worker's pump can never finalize and would bail once _proc is dropped, so if
# the row was never created (start-time create failed) the escalation creates it and
# finalizes by id itself, recording the terminal state before dropping the handle.
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id = "job_s"
b._db_config = {"model_name": "m"} # so _ensure_db_run_created can create the row
b._db_run_created = False # start-time create failed
b._proc = _FakeProc(alive = True) # wedged: still reports alive
b._should_stop = True
b._progress.is_training = True
b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_s")
assert [c["id"] for c in recs["created"]] == ["job_s"], "must create the missing row"
assert [f["id"] for f in recs["finished"]] == ["job_s"], "must finish the created row by id"
assert b._proc is None, "handle dropped only after the terminal state is recorded"
assert b._db_run_created is True
def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch):
# If a run replaces the worker while the finalize DB write is in flight, the final _proc
# drop must leave the new run's handle intact (re-guarded on target_proc).
b = TrainingBackend()
b.current_job_id = "job_old"
b._db_run_created = True
old_proc = _FakeProc(alive = False)
new_proc = _FakeProc(alive = True)
b._proc = old_proc
def hijack(*a):
b._proc = new_proc # a new run takes over during the finalize
monkeypatch.setattr(b, "_finish_stopped_run", hijack)
b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old")
assert b._proc is new_proc, "must not drop the handle a new run installed during finalize"
def _make_finish_raise(monkeypatch, calls):
fn = sys.modules["storage.studio_db"]
def _boom(**kw):
calls.append(kw)
raise RuntimeError("database is locked")
fn.finish_run = _boom
def test_finish_stopped_run_retries_then_unclaims_on_db_error(monkeypatch):
# The watchdog is the sole finalizer once _proc is dropped, so a transient DB error is
# retried a few times; on final failure the finalize is unclaimed (run still current).
monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0)
_install_fake_db(monkeypatch)
tries: list = []
_make_finish_raise(monkeypatch, tries)
b = TrainingBackend()
b.current_job_id = "job_r"
b._run_finalized = True # the caller (escalation) already claimed
b._finish_stopped_run("job_r", None, [{"step": 1}], 1, None, None, [])
assert len(tries) == 3, "a transient DB error must be retried before giving up"
assert b._run_finalized is False, "a persistent DB error must unclaim the finalize"
def test_finish_stopped_run_error_leaves_new_run_untouched(monkeypatch):
# If the watched run was superseded, a DB error must not unclaim the new run's finalize.
monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0)
_install_fake_db(monkeypatch)
_make_finish_raise(monkeypatch, [])
b = TrainingBackend()
b.current_job_id = "job_new" # a new run is live
b._run_finalized = True # the new run's flag
b._finish_stopped_run("job_old", None, [{"step": 1}], 1, None, None, [])
assert b._run_finalized is True, "must not unclaim the new run's finalize"

File diff suppressed because it is too large Load diff

View file

@ -2550,3 +2550,649 @@ class TestHfEndpointUnreachable:
t0 = time.time()
result = hf_endpoint_unreachable(timeout = 2)
assert result is True and (time.time() - t0) < 6.0
class TestLatestTierActiveFor:
"""latest_tier_active_for: the 16-bit guard for the consented latest sidecar."""
@staticmethod
def _pin(
monkeypatch,
tv,
version = "5.13.1",
):
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: version)
monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: None)
def test_true_when_tier_latest(self, monkeypatch):
import utils.transformers_version as tv
self._pin(monkeypatch, tv)
monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest")
assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is True
def test_false_for_fixed_tiers(self, monkeypatch):
import utils.transformers_version as tv
self._pin(monkeypatch, tv)
for tier in ("default", "530", "550", "510"):
monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, _t = tier, **k: _t)
assert tv.latest_tier_active_for("some/model") is False
def test_false_without_pin_and_no_resolution(self, monkeypatch):
"""No sidecar pin returns False before any tier or network resolution."""
import utils.transformers_version as tv
def _boom(*a, **k):
raise AssertionError("must not resolve without a pin")
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None)
monkeypatch.setattr(tv, "_remote_lora_base", _boom)
monkeypatch.setattr(tv, "get_transformers_tier", _boom)
assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is False
def test_never_raises(self, monkeypatch):
import utils.transformers_version as tv
def _boom(*a, **k):
raise RuntimeError("tier resolution exploded")
self._pin(monkeypatch, tv)
monkeypatch.setattr(tv, "get_transformers_tier", _boom)
assert tv.latest_tier_active_for("some/model") is False
def test_remote_lora_base_is_resolved(self, monkeypatch):
"""A remote adapter is judged by its base model, like worker activation."""
import utils.transformers_version as tv
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.13.1")
monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: "Zyphra/ZAYA1-8B")
tiers = {"Zyphra/ZAYA1-8B": "latest"}
monkeypatch.setattr(
tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default")
)
assert tv.latest_tier_active_for("someuser/zaya-lora") is True
def test_local_checkpoint_config_upgrades(self, monkeypatch, tmp_path):
"""An adapter dir with its own config.json merges tiers like activation does."""
import utils.transformers_version as tv
adapter = tmp_path / "ckpt"
adapter.mkdir()
(adapter / "adapter_config.json").write_text("{}")
(adapter / "adapter_model.safetensors").write_text("x")
(adapter / "config.json").write_text("{}")
self._pin(monkeypatch, tv)
monkeypatch.setattr(tv, "_resolve_base_model", lambda name: "base/model")
tiers = {"base/model": "default", str(adapter): "latest"}
monkeypatch.setattr(
tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default")
)
assert tv.latest_tier_active_for(str(adapter)) is True
class TestLatestTierForces16Bit:
"""The inference worker and load route refuse bnb 4-bit on the latest sidecar."""
def _read(self, rel):
backend_dir = Path(__file__).resolve().parent.parent
return (backend_dir / rel).read_text()
def test_worker_guard_present(self):
src = self._read("core/inference/worker.py")
assert "latest_tier_active_for" in src, (
"core/inference/worker.py must force load_in_4bit=False when "
"latest_tier_active_for(model) is true: transformers' grouped-MoE "
"kernels crash on bnb-quantized expert weights for brand-new "
"architectures."
)
def test_route_guard_present(self):
src = self._read("routes/inference.py")
assert "latest_tier_active_for" in src, (
"routes/inference.py must size the VRAM guard with the same 16-bit "
"flip the worker applies for latest-sidecar models."
)
def test_validate_route_mirrors_16bit_flip(self):
# Without the same flip, /validate sizes 4-bit and /load then 409s.
src = self._read("routes/inference.py")
body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0]
assert "latest_tier_active_for" in body, (
"validate_model must apply the latest-sidecar 16-bit flip before "
"_guard_chat_load_against_training so /validate and /load agree."
)
# First-time loads have no pin yet, so an installable upgrade must also size 16-bit.
assert body.index("check_upgrade_for_model") < body.index(
"_guard_chat_load_against_training"
), "the upgrade check must run before the training guard"
assert (
"supported_in_pypi" in body.split("_guard_chat_load_against_training")[0]
), "an installable upgrade must force 16-bit sizing for the guard"
def test_validate_offered_upgrade_preserves_custom_code_4bit(self):
# A merely-offered (not installed) upgrade must NOT force 16-bit sizing when the
# model has a custom-code (auto_map) fallback: /load loads it 4-bit without the
# install, and the install route refuses during active training, so 16-bit sizing
# here would 409 the only viable 4-bit path.
src = self._read("routes/inference.py")
body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0]
flip = body.split("Mirror /load's latest-sidecar 16-bit flip", 1)[1].split(
"_guard_chat_load_against_training", 1
)[0]
assert "not requires_trust_remote_code" in flip, (
"the offered-upgrade 16-bit flip must be gated on the absence of a custom-code "
"fallback so /validate does not 409 a 4-bit load /load would allow"
)
# requires_trust_remote_code must be resolved before the flip consumes it.
assert body.index("requires_trust_remote_code = any(") < body.index(
"not requires_trust_remote_code"
)
def test_install_route_guards_active_latest_workers(self):
# Stage-and-swap replaces .venv_t5_latest in place, so a live worker on the
# old sidecar would lazy-import files from the new version.
src = self._read("routes/inference.py")
body = src.split("async def install_latest_transformers_route", 1)[1].split(
"\nasync def ", 1
)[0]
assert (
"is_training_active" in body
and "is_export_active" in body
and "inference_lifecycle_gate" in body
), (
"install_latest_transformers_route must refuse while training or export "
"runs, and hold the lifecycle gate while unloading the chat model and "
"swapping the sidecar."
)
# The unload (via before_swap so failed installs keep the model), the export-worker
# teardown, and the install must all sit INSIDE the gate so no /load interleaves.
assert "unload_model(active)" in body
assert "cleanup_memory()" in body
# Export teardown precedes the chat unload so its failure aborts with the model still loaded.
assert body.index("cleanup_memory()") < body.index("unload_model(active)")
assert "install_latest_transformers(" in body and "_unload_before_swap" in body
# The gate must be owned by the shielded task, not the request coroutine: a cancelled
# POST unwinding an async-with would release the only guard /load honors mid-install.
gated_task = body.split("async def _gated_install", 1)[1]
assert "inference_lifecycle_gate():" in gated_task
assert "asyncio.to_thread(_run_install)" in gated_task
# The reservation must be taken BEFORE the (awaitable) gate wait, or a
# training/export start could slip in while this request queues on the gate.
assert body.index("try_begin_sidecar_swap()") < body.index(
"inference_lifecycle_gate():"
), "the swap reservation must be raised before waiting on the lifecycle gate"
# A failed teardown must abort the swap (raise), not fall through to it.
assert body.count("raise RuntimeError") >= 3, (
"export, chat-unload, and idle-worker teardown failures must raise so "
"the staged install never swaps under a live worker"
)
# The installer thread owns (and releases) the reservation, shielded from
# request cancellation, so a cancelled POST cannot unlock a live swap.
assert "asyncio.shield" in body and "end_sidecar_swap()" in body
# In-flight generation streams predate the gate; the route refuses rather than kill them
# via the before_swap unload. The count is rechecked UNDER the gate, since a wait on a
# long /load outlasts the pre-gate fast path and streams take this same gate.
assert "other_inference_request_count" in body
gated_task = body.split("async def _gated_install", 1)[1]
assert "other_inference_request_count" in gated_task
def test_start_routes_refuse_during_install(self):
# A worker spawned mid-swap could activate a half-replaced sidecar.
training = self._read("routes/training.py")
start = training.split("async def start_training", 1)[1].split("\nasync def ", 1)[0]
assert (
"is_install_in_progress" in start
), "training /start must refuse while a transformers install is in progress"
export = self._read("routes/export.py")
helper = export.split("def _ensure_export_supported", 1)[1].split("\ndef ", 1)[0]
assert (
"is_install_in_progress" in helper
), "mutating export routes must refuse while a transformers install is in progress"
def test_spawn_sites_recheck_reservation(self):
# The route-level guards are one-shot; validation between them and the
# actual spawn can outlast an install's start, so the spawn itself rechecks.
training = self._read("core/training/training.py")
assert (
training.count("sidecar_swap_in_progress()") >= 2
), "both training spawn sites must recheck the sidecar swap reservation"
export = self._read("core/export/orchestrator.py")
spawn = export.split("def _spawn_subprocess", 1)[1].split("\n def ", 1)[0]
assert (
"sidecar_swap_kind()" in spawn
), "the export subprocess spawn must recheck the sidecar swap reservation"
# Training marks the spawn active BEFORE its recheck, so either side sees the other:
# is_training_active covers the window between proc.start() and the _proc assignment.
assert training.index("self._spawn_in_progress = True") < training.index(
"if sidecar_swap_in_progress():"
)
active = training.split("def is_training_active", 1)[1].split("\n def ", 1)[0]
assert "_spawn_in_progress" in active
# Export load-checkpoint refuses BEFORE tearing down the old worker, so a
# lost race against an install keeps the loaded checkpoint (no bare 500).
loadck = export.split("def load_checkpoint", 1)[1].split("\n def ", 1)[0]
assert loadck.index("sidecar_swap_in_progress()") < loadck.index("_shutdown_subprocess()")
# The training handshake precedes the VRAM-freeing before_spawn hook, so
# losing the race never tears down chat/export for a run that won't spawn.
assert training.index("self._spawn_in_progress = True") < training.index("before_spawn()")
# The spawn-time export check is op-aware for installs (the install side
# aborts on is_export_active) but always refuses for repairs, which have
# no such abort and can be rebuilding the sidecar right now.
assert (
'_swap_kind == "repair" or (_swap_kind is not None and not self._export_active)'
in spawn
)
class TestSidecarSwapReservation:
"""The lazy repair takes the same reservation the install route and worker starts use."""
def _repair_setup(self, monkeypatch, tmp_path):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
monkeypatch.setattr(
tv,
"_latest_pin_data",
lambda: {
"version": "5.99.0",
"packages": ["transformers==5.99.0"],
},
)
monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: False)
monkeypatch.setattr(tv, "_env_offline", lambda: False)
return tv
def test_repair_holds_reservation_during_swap(self, monkeypatch, tmp_path):
tv = self._repair_setup(monkeypatch, tmp_path)
seen = {}
def _fake_swap(
version,
packages,
before_swap = None,
):
seen["active_during_swap"] = tv.sidecar_swap_in_progress()
return True
monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _fake_swap)
assert tv._ensure_venv_t5_latest_exists() is True
assert seen["active_during_swap"] is True
assert tv.sidecar_swap_in_progress() is False
def test_foreign_process_lock_file_visible(self, monkeypatch, tmp_path):
"""A repair in a LIVE worker subprocess is seen (via the lock file) by this
process, and its lock is never broken while the owner is alive."""
import os
import time
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
lock = tv._swap_lock_path()
lock.parent.mkdir(parents = True, exist_ok = True)
# A live owner (this process): visible and never reclaimed, even once aged past
# the cutoff -- a slow but live pip install must keep its lock.
lock.write_text('{"pid": %d}' % os.getpid())
assert tv.sidecar_swap_in_progress() is True
assert tv.try_begin_sidecar_swap() is False
old_ts = time.time() - 3 * 60 * 60
os.utime(lock, (old_ts, old_ts))
assert tv.sidecar_swap_in_progress() is True
assert tv.try_begin_sidecar_swap() is False
def test_dead_owner_lock_reclaimed_promptly(self, monkeypatch, tmp_path):
"""A fresh lock whose recorded owner is dead is reclaimed at once, not after the
long cutoff: a crash mid-install must not wedge loads/training/export for hours."""
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
lock = tv._swap_lock_path()
lock.parent.mkdir(parents = True, exist_ok = True)
# 999999 is not a live PID: a fresh dead-owner lock is immediately stale.
lock.write_text('{"pid": 999999, "kind": "install"}')
assert tv._pid_alive(999999) is False
assert tv.sidecar_swap_in_progress() is False
assert tv.try_begin_sidecar_swap() is True
try:
assert lock.is_file()
finally:
tv.end_sidecar_swap()
assert not lock.exists()
def test_unreadable_pid_lock_uses_age_cutoff(self, monkeypatch, tmp_path):
"""A lock with no readable owner PID (mid create-before-write, or corrupt) is not
reclaimed while fresh -- only after the long cutoff -- so a lock a live owner just
created is not stolen before its PID lands."""
import os
import time
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
lock = tv._swap_lock_path()
lock.parent.mkdir(parents = True, exist_ok = True)
lock.write_text("") # created but metadata not yet written
assert tv.sidecar_swap_in_progress() is True
old_ts = time.time() - (tv._SWAP_LOCK_STALE_SECS + 60)
os.utime(lock, (old_ts, old_ts))
assert tv.sidecar_swap_in_progress() is False
def test_repair_refused_while_install_holds_reservation(self, monkeypatch, tmp_path):
tv = self._repair_setup(monkeypatch, tmp_path)
def _must_not_run(*a, **k):
raise AssertionError("repair must not swap while an install is in progress")
monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _must_not_run)
assert tv.try_begin_sidecar_swap() is True
try:
assert tv._ensure_venv_t5_latest_exists() is False
finally:
tv.end_sidecar_swap()
class TestRecoverStrandedSidecar:
"""A swap whose activation rename AND rollback both fail strands the previous sidecar
at .old with no live dir (its pin marker went with it). Reading the pin self-heals it,
but never while a swap legitimately holds the reservation."""
def _setup(self, monkeypatch, tmp_path):
import utils.transformers_version as tv
live = str(tmp_path / "venv_t5_latest")
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", live)
# Stranded state: live gone, previous sidecar (with its marker) sits at .old.
retired = Path(live + ".old")
retired.mkdir(parents = True)
(retired / tv._LATEST_PIN_MARKER).write_text(
'{"version": "5.99.0", "packages": ["transformers==5.99.0"]}'
)
return tv, Path(live), retired
def test_stranded_old_recovered_on_pin_read(self, monkeypatch, tmp_path):
tv, live, retired = self._setup(monkeypatch, tmp_path)
data = tv._latest_pin_data()
assert live.is_dir()
assert not retired.exists()
assert data is not None and data["version"] == "5.99.0"
def test_stranded_recovery_skipped_during_swap(self, monkeypatch, tmp_path):
tv, live, retired = self._setup(monkeypatch, tmp_path)
assert tv.try_begin_sidecar_swap() is True
try:
# A swap holds the reservation and may be mid-rename; do not race it.
assert tv._latest_pin_data() is None
assert not live.exists()
assert retired.is_dir()
finally:
tv.end_sidecar_swap()
# Once the swap is done, the next pin read recovers the stranded sidecar.
assert tv._latest_pin_data() is not None
assert live.is_dir()
class TestCachedLatestMappingRevalidated:
"""A cached 'latest' mapping is dropped and re-resolved when the sidecar since broke
in-process, so routing self-heals instead of trusting a mapping parsed from a sidecar
that no longer exists (which would keep routing latest-only models to a broken tier)."""
def test_broken_sidecar_drops_cached_latest_mapping(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})})
monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: False)
seen = {"n": 0}
def _fake_overlay(tier):
seen["n"] += 1
return None # broken/unavailable -> empty, uncached
monkeypatch.setattr(tv, "_overlay_transformers_dir", _fake_overlay)
assert tv._config_model_types("latest") == frozenset()
assert seen["n"] == 1 # re-resolved, not served from the stale cache
assert "latest" not in tv._config_mapping_cache
def test_intact_sidecar_serves_cached_latest_mapping(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})})
monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True)
monkeypatch.setattr(
tv,
"_overlay_transformers_dir",
lambda tier: pytest.fail("intact sidecar must serve the cache without re-resolving"),
)
assert tv._config_model_types("latest") == frozenset({"brandnew"})
def test_non_latest_cache_not_revalidated(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_config_mapping_cache", {"530": frozenset({"gemma3"})})
monkeypatch.setattr(
tv,
"_latest_sidecar_intact",
lambda: pytest.fail("non-latest tiers must not pay the sidecar-intact check"),
)
assert tv._config_model_types("530") == frozenset({"gemma3"})
def test_deleted_pin_drops_cached_latest_mapping(self, monkeypatch, tmp_path):
# A pin marker deleted after the mapping was cached makes _latest_pin_data None;
# the cache must be dropped (not trusted), so routing re-resolves to no latest tier
# rather than routing to a latest tier that then fails worker activation.
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False)
monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})})
# No pin marker on disk -> _latest_pin_data() is None -> not intact.
assert tv._latest_sidecar_intact() is False
assert tv._config_model_types("latest") == frozenset()
assert "latest" not in tv._config_mapping_cache
class TestOverlayRepairsIncompleteSidecar:
"""Routing self-heals a pinned latest sidecar that is present but incomplete,
not only one whose transformers/ dir vanished: workers refuse parent-only
repairs, so a sidecar missing a pinned package would fail every load."""
def _setup(self, monkeypatch, tmp_path, valid):
import utils.transformers_version as tv
live = tmp_path / "venv_t5_latest"
(live / "transformers").mkdir(parents = True)
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live))
monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False)
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0")
monkeypatch.setattr(
tv,
"_latest_pin_data",
lambda: {"version": "5.99.0", "packages": ["transformers==5.99.0", "tiktoken"]},
)
monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: valid)
monkeypatch.setattr(tv, "_latest_repair_failed_at", 0.0)
return tv
def test_incomplete_sidecar_triggers_repair(self, monkeypatch, tmp_path):
tv = self._setup(monkeypatch, tmp_path, valid = False)
called = {"n": 0}
def _fake_repair():
called["n"] += 1
return True
monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair)
src = tv._overlay_transformers_dir("latest")
assert called["n"] == 1
assert src == str(tmp_path / "venv_t5_latest" / "transformers")
def test_intact_sidecar_skips_repair(self, monkeypatch, tmp_path):
tv = self._setup(monkeypatch, tmp_path, valid = True)
def _must_not_run():
raise AssertionError("intact sidecar must not trigger a repair")
monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _must_not_run)
assert tv._overlay_transformers_dir("latest") == str(
tmp_path / "venv_t5_latest" / "transformers"
)
def test_failed_repair_backs_off(self, monkeypatch, tmp_path):
tv = self._setup(monkeypatch, tmp_path, valid = False)
called = {"n": 0}
def _fake_repair():
called["n"] += 1
return False
monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair)
# A failed repair must not route through the broken sidecar, neither on
# the failing attempt nor while the backoff suppresses the next attempt.
assert tv._overlay_transformers_dir("latest") is None
assert tv._overlay_transformers_dir("latest") is None
assert called["n"] == 1
class TestStageAndSwapBeforeSwap:
"""before_swap fires only when the staged install succeeded and the swap is next."""
def _setup(self, monkeypatch, tmp_path, build_ok):
import utils.transformers_version as tv
live = tmp_path / "venv_latest"
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live))
def _fake_build(target, packages, label):
if build_ok:
Path(target).mkdir(parents = True, exist_ok = True)
return build_ok
monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_build)
return tv, live
def test_called_after_successful_staging(self, monkeypatch, tmp_path):
tv, live = self._setup(monkeypatch, tmp_path, build_ok = True)
calls = []
assert tv._stage_and_swap_latest_venv(
"5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1)
)
assert calls == [1] and live.is_dir()
def test_not_called_when_staging_fails(self, monkeypatch, tmp_path):
tv, live = self._setup(monkeypatch, tmp_path, build_ok = False)
calls = []
assert not tv._stage_and_swap_latest_venv(
"5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1)
)
assert calls == [] and not live.exists()
def test_failure_in_before_swap_keeps_previous_sidecar(self, monkeypatch, tmp_path):
tv, live = self._setup(monkeypatch, tmp_path, build_ok = True)
live.mkdir()
(live / "sentinel").write_text("old")
def _boom():
raise RuntimeError("worker teardown failed")
assert not tv._stage_and_swap_latest_venv(
"5.99.0", ("transformers==5.99.0",), before_swap = _boom
)
assert (live / "sentinel").read_text() == "old"
class TestKillSwitchBeatsMappingCache:
def test_cached_latest_probe_ignored_when_disabled(self, monkeypatch):
import utils.transformers_version as tv
key = tv._probe_cache_key("some/model")
monkeypatch.setitem(tv._probe_tier_cache, key, "latest")
monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1")
# With the switch set, the cached latest entry must not short-circuit;
# the probe re-resolves against the non-latest order (stub it to 530).
monkeypatch.setattr(tv, "_probe_tier_venvs", lambda: {})
monkeypatch.setattr(tv, "_probe_tier_order", lambda: ())
assert tv._probe_tier("some/model", None, "test") != "latest"
# Cached non-latest entries and the unset switch still short-circuit.
monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS")
assert tv._probe_tier("some/model", None, "test") == "latest"
def test_cached_latest_mapping_ignored_when_disabled(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setitem(tv._config_mapping_cache, "latest", frozenset({"brandnew"}))
# The cache is trusted only when the sidecar is intact; hold it intact so this
# test isolates the kill switch, not the sidecar-revalidation path.
monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True)
monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1")
assert tv._config_model_types("latest") == frozenset()
monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS")
assert tv._config_model_types("latest") == frozenset({"brandnew"})
class TestRaiseTierForNested:
"""_raise_tier_for_nested: a wrapper's nested model_type can raise a fast-path tier."""
def _patch_types(self, monkeypatch, per_tier):
import utils.transformers_version as tv
monkeypatch.setattr(
tv, "_config_model_types", lambda tier: frozenset(per_tier.get(tier, ()))
)
def test_nested_latest_only_type_raises(self, monkeypatch):
import utils.transformers_version as tv
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}})
cfg = {"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}}
assert tv._raise_tier_for_nested(cfg, "550") == "latest"
def test_never_lowers_a_fast_path_tier(self, monkeypatch):
import utils.transformers_version as tv
# Mapping alone would say 530, but the fast path (e.g. a name override) said 550.
self._patch_types(monkeypatch, {"530": {"qwen3_5"}, "550": {"qwen3_5"}})
assert tv._raise_tier_for_nested({"model_type": "qwen3_5"}, "550") == "550"
def test_no_config_keeps_tier(self):
import utils.transformers_version as tv
assert tv._raise_tier_for_nested(None, "550") == "550"
def test_unknown_nested_type_never_vetoes(self, monkeypatch):
import utils.transformers_version as tv
# A nested type unknown everywhere (not even latest) keeps the fast path.
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4"}})
cfg = {"model_type": "gemma4", "text_config": {"model_type": "unreleased"}}
assert tv._raise_tier_for_nested(cfg, "550") == "550"
def test_name_fast_path_folds_when_latest_pinned(self, monkeypatch):
"""A fixed-tier name match with a latest-only model_type routes to latest
once the sidecar is pinned; without a pin the name tier stands (no I/O)."""
import utils.transformers_version as tv
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"brandnew_arch"}})
monkeypatch.setattr(tv, "_tier_from_name", lambda name: ("550", "gemma-4"))
monkeypatch.setattr(
tv, "_load_config_json", lambda name, tok = None: {"model_type": "brandnew_arch"}
)
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0")
assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "latest"
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None)
monkeypatch.setattr(
tv,
"_load_config_json",
lambda name, tok = None: (_ for _ in ()).throw(AssertionError("no I/O without a pin")),
)
assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "550"
def test_fast_path_folds_nested_tier(self, monkeypatch, tmp_path):
"""End to end: a local wrapper config on a fixed fast path routes to latest
when its nested type only exists in the installed latest sidecar."""
import utils.transformers_version as tv
ckpt = tmp_path / "wrapper"
ckpt.mkdir()
(ckpt / "config.json").write_text(
json.dumps({"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}})
)
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}})
monkeypatch.setattr(tv, "_config_needs_510", lambda cfg: False)
monkeypatch.setattr(tv, "_config_needs_550", lambda cfg: True)
assert tv.get_transformers_tier(str(ckpt), probe = False) == "latest"

View file

@ -0,0 +1,354 @@
# 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 ast
import os
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Optional
from utils.paths import external_media
_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_routes_function(name: str, ns_extra: Optional[dict] = None) -> dict:
"""Exec one top-level function from routes/models.py without importing the module (which pulls in FastAPI)."""
tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name)
module = ast.Module(body = [fn], type_ignores = [])
ast.fix_missing_locations(module)
ns = {"os": os, "Path": Path, "Optional": Optional}
if ns_extra:
ns.update(ns_extra)
exec(compile(module, "<extracted routes/models.py>", "exec"), ns)
return ns
def _stub_windows(monkeypatch, existing_drives):
"""Simulate Windows exposing only *existing_drives* (e.g. {"C", "D"}) as readable roots, independent of the host FS.
Overriding _active_windows_drive_bitmask keeps it deterministic even on a
real Windows host, where live GetLogicalDrives would return the actual layout."""
monkeypatch.setattr(external_media.platform, "system", lambda: "Windows")
mask = sum(1 << (ord(d.upper()) - ord("A")) for d in existing_drives)
monkeypatch.setattr(external_media, "_active_windows_drive_bitmask", lambda: mask)
present = {f"{d.upper()}:\\" for d in existing_drives}
monkeypatch.setattr(external_media.os.path, "isdir", lambda p: str(p) in present)
monkeypatch.setattr(external_media.os, "access", lambda p, _mode: str(p) in present)
def test_windows_drive_roots_empty_off_windows(monkeypatch):
# Regression guard: the helper is a no-op on Linux/macOS so it can't change the allowlist on the platforms CI runs on.
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
assert external_media.windows_drive_roots() == []
monkeypatch.setattr(external_media.platform, "system", lambda: "Darwin")
assert external_media.windows_drive_roots() == []
def test_windows_drive_roots_lists_readable_drives(monkeypatch):
_stub_windows(monkeypatch, {"C", "D", "E"})
roots = external_media.windows_drive_roots(drive_letters = "CDEF")
# F is absent, so it is skipped; the rest are exposed in order.
assert roots == [Path("C:\\"), Path("D:\\"), Path("E:\\")]
def test_windows_drive_roots_skips_absent_and_unreadable(monkeypatch):
_stub_windows(monkeypatch, {"C"})
roots = external_media.windows_drive_roots(drive_letters = "CDE")
assert roots == [Path("C:\\")]
def test_windows_drive_roots_ignores_bad_letters_and_dedupes(monkeypatch):
_stub_windows(monkeypatch, {"C", "D"})
roots = external_media.windows_drive_roots(
drive_letters = ["c:", "C", "D", "1", "AB", "", " d "],
)
assert roots == [Path("C:\\"), Path("D:\\")]
def test_readable_dir_within_times_out(monkeypatch):
# A probe that outlives the timeout is reported not-readable, so a hung
# (disconnected mapped network) drive is skipped instead of blocking.
import time
monkeypatch.setattr(external_media.os.path, "isdir", lambda p: time.sleep(5) or True)
monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True)
start = time.monotonic()
ok = external_media._readable_dir_within("Z:\\", timeout = 0.2)
elapsed = time.monotonic() - start
assert ok is False
assert elapsed < 3.0 # returned on the timeout, did not wait out the 5s stall
def test_readable_dir_within_reports_fast_probe(monkeypatch):
monkeypatch.setattr(external_media.os.path, "isdir", lambda p: True)
monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True)
assert external_media._readable_dir_within("C:\\", timeout = 2.0) is True
def test_windows_drive_roots_skips_hung_drive(monkeypatch):
# A disconnected mapped drive stays set in the bitmask and its os.path.isdir
# stalls; it must be skipped without stalling enumeration. C answers, D hangs,
# so only C is listed, bounded by the per-drive timeout, not the stall.
import time
monkeypatch.setattr(external_media.platform, "system", lambda: "Windows")
monkeypatch.setattr(
external_media,
"_active_windows_drive_bitmask",
lambda: sum(1 << (ord(d) - ord("A")) for d in "CD"),
)
monkeypatch.setattr(external_media, "_DRIVE_PROBE_TIMEOUT_S", 0.2)
def _isdir(p):
if str(p) == "D:\\":
time.sleep(5) # simulate the reconnect stall
return True
return str(p) == "C:\\"
monkeypatch.setattr(external_media.os.path, "isdir", _isdir)
monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True)
start = time.monotonic()
roots = external_media.windows_drive_roots(drive_letters = "CD")
elapsed = time.monotonic() - start
assert roots == [Path("C:\\")]
assert elapsed < 3.0 # bounded by the per-drive timeout, not the 5s stall
def test_windows_drive_roots_probes_hung_drives_in_parallel(monkeypatch):
# Several disconnected mapped drives must add ~one timeout total, not one
# per drive: C answers fast, D/E/F stall. The concurrent probe stays bounded
# by a single deadline where serial probing would cost ~4x the timeout.
import time
monkeypatch.setattr(external_media.platform, "system", lambda: "Windows")
monkeypatch.setattr(
external_media,
"_active_windows_drive_bitmask",
lambda: sum(1 << (ord(d) - ord("A")) for d in "CDEF"),
)
timeout = 0.2
monkeypatch.setattr(external_media, "_DRIVE_PROBE_TIMEOUT_S", timeout)
def _isdir(p):
if str(p) == "C:\\":
return True
time.sleep(5) # every other drive simulates a reconnect stall
return True
monkeypatch.setattr(external_media.os.path, "isdir", _isdir)
monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True)
start = time.monotonic()
roots = external_media.windows_drive_roots(drive_letters = "CDEF")
elapsed = time.monotonic() - start
assert roots == [Path("C:\\")]
# 3 stalled drives probed in parallel finish within ~1 timeout, well under the ~3*timeout a serial probe would take.
assert elapsed < 3 * timeout
def test_browse_allowlist_includes_windows_drive_roots(monkeypatch, tmp_path):
# End-to-end wiring: windows_drive_roots() output flows into the browse
# allowlist built by routes/models.py, mirroring the Linux media-mounts test.
tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
function_names = {
"_build_browse_allowlist",
"_browse_relative_parts",
"_is_path_inside_allowlist",
"_match_browse_child",
"_normalize_browse_request_path",
"_resolve_browse_target",
}
functions = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name in function_names
]
module = ast.Module(body = functions, type_ignores = [])
ast.fix_missing_locations(module)
home = tmp_path / "home"
drive_root = tmp_path / "D_drive"
model_dir = drive_root / "modelsAI" / "gguf"
home.mkdir()
model_dir.mkdir(parents = True)
fake_paths = SimpleNamespace(
hf_default_cache_dir = lambda: tmp_path / "missing-default-hf",
legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf",
well_known_model_dirs = lambda: [],
studio_root = lambda: tmp_path / "missing-studio",
outputs_root = lambda: tmp_path / "missing-outputs",
exports_root = lambda: tmp_path / "missing-exports",
)
fake_external_media = SimpleNamespace(
linux_run_media_mount_roots = lambda: [],
windows_drive_roots = lambda: [drive_root],
)
fake_studio_db = SimpleNamespace(
list_scan_folders = lambda: [],
contains_sensitive_path_component = lambda _p: False,
# The simulated D:\ root maps to a tmp_path dir, not a denied system path.
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)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db)
ns = {
"HTTPException": _HTTPException,
"os": os,
"Path": Path,
"Optional": Optional,
"_safe_is_dir": lambda p: Path(p).is_dir(),
"_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf",
"logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None),
}
exec(compile(module, "<extracted routes/models.py>", "exec"), ns)
allowlist = ns["_build_browse_allowlist"]()
# The simulated Windows drive root is now browsable, and a model dir on it resolves.
assert drive_root.resolve() in allowlist
assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve()
def test_build_browse_allowlist_reuses_passed_roots(monkeypatch, tmp_path):
# Double-probe fix: a browse request probes the drive/media roots once and
# passes them in, so _build_browse_allowlist must NOT scan
# windows_drive_roots() again (a disconnected drive would double the stall).
tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
functions = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "_build_browse_allowlist"
]
module = ast.Module(body = functions, type_ignores = [])
ast.fix_missing_locations(module)
drive_root = tmp_path / "D_drive"
drive_root.mkdir()
calls = {"drive": 0, "media": 0}
def _drive_roots():
calls["drive"] += 1
return [drive_root]
def _media_roots():
calls["media"] += 1
return []
fake_paths = SimpleNamespace(
hf_default_cache_dir = lambda: tmp_path / "missing-default-hf",
legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf",
well_known_model_dirs = lambda: [],
studio_root = lambda: tmp_path / "missing-studio",
outputs_root = lambda: tmp_path / "missing-outputs",
exports_root = lambda: tmp_path / "missing-exports",
)
fake_external_media = SimpleNamespace(
linux_run_media_mount_roots = _media_roots,
windows_drive_roots = _drive_roots,
)
fake_studio_db = SimpleNamespace(list_scan_folders = lambda: [])
monkeypatch.setitem(sys.modules, "utils.paths", fake_paths)
monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db)
ns = {
"os": os,
"Path": Path,
"Optional": Optional,
"_safe_is_dir": lambda p: Path(p).is_dir(),
"_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf",
"logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None),
}
exec(compile(module, "<extracted routes/models.py>", "exec"), ns)
build = ns["_build_browse_allowlist"]
# Roots passed in -> neither helper is probed, but the roots still flow in.
allowlist = build([], [drive_root])
assert calls == {"drive": 0, "media": 0}
assert drive_root.resolve() in allowlist
# No args -> each helper is probed exactly once.
build()
assert calls == {"drive": 1, "media": 1}
def test_is_path_inside_allowlist_real_descendants_and_siblings(tmp_path):
# Component-wise containment (commonpath): a genuine descendant is allowed,
# but a sibling sharing only a string prefix ("models_root_evil" vs
# "models_root") is not, which the old startswith check could miss.
ns = _extract_routes_function("_is_path_inside_allowlist")
root = tmp_path / "models_root"
child = root / "gguf" / "qwen"
sibling = tmp_path / "models_root_evil"
child.mkdir(parents = True)
sibling.mkdir()
is_inside = ns["_is_path_inside_allowlist"]
assert is_inside(root, [root]) is True # the root itself
assert is_inside(child, [root]) is True # a genuine descendant
assert is_inside(sibling, [root]) is False # prefix-collision sibling
def test_is_path_inside_allowlist_posix_root_does_not_authorize_descendants(monkeypatch):
# Regression for the reported POSIX "/" unlock: a bare filesystem root may
# match itself but must NOT authorize arbitrary descendants such as /etc.
ns = _extract_routes_function("_is_path_inside_allowlist")
monkeypatch.setattr(os.path, "realpath", lambda p: str(p)) # keep "/" intact
is_inside = ns["_is_path_inside_allowlist"]
assert is_inside("/", ["/"]) is True # the root itself
assert is_inside("/etc", ["/"]) is False # not a licensed descendant
assert is_inside("/root/models", ["/"]) is False
def test_is_path_inside_allowlist_windows_drive_root_descendants():
# Exercise the Windows drive-root branch on a POSIX host by backing os.path
# with ntpath and an identity realpath (the simulated drives don't exist
# here). A drive root authorizes its descendants; a different drive does not.
import ntpath
win_os = SimpleNamespace(
sep = "\\",
path = SimpleNamespace(
normcase = ntpath.normcase,
realpath = lambda p: str(p),
splitdrive = ntpath.splitdrive,
dirname = ntpath.dirname,
commonpath = ntpath.commonpath,
),
)
ns = _extract_routes_function("_is_path_inside_allowlist", {"os": win_os})
is_inside = ns["_is_path_inside_allowlist"]
assert is_inside("D:\\", ["D:\\"]) is True # drive root itself
assert is_inside("D:\\models", ["D:\\"]) is True # descendant on the drive
assert is_inside("D:\\models\\gguf", ["D:\\"]) is True # deeper descendant
assert is_inside("d:\\models", ["D:\\"]) is True # case-insensitive drive letter
assert is_inside("C:\\Users", ["D:\\"]) is False # different drive
assert is_inside("D:\\models", ["E:\\"]) is False

View file

@ -698,6 +698,25 @@ if backend_dir not in sys.path:
try:
from transformers import AutoConfig
# Union the ACTIVE sidecar's registry into the inlined parent-process sets
# so architectures only the sidecar knows still classify correctly.
try:
from transformers.models.auto import modeling_auto as _ma
for _attr in ("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES",
"MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES"):
_d = dict(getattr(_ma, _attr, None) or {})
_VLM_MODEL_TYPES |= set(_d)
_VLM_CLASS_NAMES |= set(_d.values())
for _attr in ("MODEL_FOR_CTC_MAPPING_NAMES",
"MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES",
"MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES",
"MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES"):
_AUDIO_ONLY_MODEL_TYPES |= set(dict(getattr(_ma, _attr, None) or {}))
except Exception:
pass
# Capability detection never executes model repo code.
kwargs = {"trust_remote_code": False}
if token:
@ -727,13 +746,23 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
"""
token_arg = hf_token or ""
# Latest-only architectures need the latest sidecar for AutoConfig;
# other tiers keep the 5.5 sidecar.
sidecar_dir = _VENV_T5_DIR
try:
from utils.transformers_version import _VENV_T5_LATEST_DIR, get_transformers_tier
if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
sidecar_dir = _VENV_T5_LATEST_DIR
except Exception:
pass
try:
result = subprocess.run(
[
sys.executable,
"-c",
_VISION_CHECK_SCRIPT,
_VENV_T5_DIR,
sidecar_dir,
_BACKEND_DIR,
model_name,
token_arg,
@ -876,6 +905,17 @@ def _is_vision_model_uncached(
model_name, hf_token = hf_token, local_files_only = local_files_only
)
if raw is not None:
if raw is False and not local_files_only:
# Raw heuristics predate latest-only architectures; on the latest tier,
# trust that sidecar's AutoConfig probe over the heuristic False. An
# inconclusive probe (sidecar mid-repair, timeout) is transient: return
# None so the heuristic False is not cached and the model is re-probed.
try:
from utils.transformers_version import get_transformers_tier
if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
except Exception:
pass
return raw
# Raw read failed transiently: fall back to AutoConfig (remote code DISABLED), via a

View file

@ -8,6 +8,10 @@ from __future__ import annotations
import getpass
import os
import platform
import string
import threading
import time
from collections.abc import Iterable
from pathlib import Path
from utils.paths.sensitive import (
@ -16,6 +20,33 @@ from utils.paths.sensitive import (
)
def is_local_filesystem_root(path: str, *, _pathmod = os.path) -> bool:
"""True for a bare local filesystem root -- POSIX ``/``, a drive root ``C:\\``,
or a device-namespace volume root like ``\\\\?\\C:\\`` or
``\\\\?\\Volume{GUID}\\`` -- which sit above denied system dirs, but NOT a UNC
share root (``\\\\server\\share`` or its ``\\\\?\\UNC\\...`` form), which has
none under it and was registerable before this guard. ``splitdrive`` is empty
on POSIX servers, so this reduces to the plain ``dirname == self`` test there.
``_pathmod`` lets tests drive ``ntpath`` semantics on a POSIX CI.
"""
# Resolve the Windows device / extended-length namespace, where \\?\C:\,
# \\.\C:\ and \\?\Volume{GUID}\ are all bare LOCAL volume roots (rejected)
# while only \\?\UNC\server\share is a UNC share (handled like \\server\share).
if path[:4].lower() in ("\\\\?\\", "\\\\.\\"):
rest = path[4:]
if rest[:4].lower() == "unc\\":
path = "\\\\" + rest[4:]
else:
# A device volume root is just the volume specifier (C:, Volume{GUID})
# with no further component; a deeper path is an ordinary folder.
core = rest.rstrip("\\/")
return "\\" not in core and "/" not in core
if _pathmod.dirname(path) != path:
return False
drive, _ = _pathmod.splitdrive(path)
return drive[:2] not in ("\\\\", "//")
def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool:
normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path)))
root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root))))
@ -98,3 +129,101 @@ def linux_run_media_mount_roots(
seen.add(key)
roots.append(resolved)
return roots
def _active_windows_drive_bitmask() -> int:
"""Active-logical-drive bitmask from ``GetLogicalDrives`` (bit 0 = ``A:``), or ``0`` when unavailable.
A fast non-blocking call that lets :func:`windows_drive_roots` skip the
``os.path.isdir`` probe on unmapped letters. A disconnected network mapping
stays set here, so it does not guard the reconnect stall on its own;
:func:`windows_drive_roots` bounds each surviving probe too. Returns ``0``
(probe every letter) when ctypes/``windll`` is missing.
"""
try:
import ctypes
return int(ctypes.windll.kernel32.GetLogicalDrives())
except Exception: # noqa: BLE001 -- best-effort; fall back to probing all letters
return 0
# A disconnected mapped drive stays set in the GetLogicalDrives bitmask, so
# ``os.path.isdir`` on it can block for tens of seconds. Bound each drive probe
# so one stale mapping cannot stall a whole folder-browser request.
_DRIVE_PROBE_TIMEOUT_S = 2.0
def _readable_dirs_within(paths: Iterable[str], timeout: float) -> set[str]:
"""Which of *paths* are readable directories, probed concurrently under one overall *timeout* (seconds).
Each path is checked (``os.path.isdir`` + ``os.access(R_OK)``) in its own
daemon thread and the call waits at most *timeout* total, not per path, so N
stalled network drives add ~timeout instead of N*timeout. A path not
answering ``True`` by the deadline is treated as unreadable. The daemon
threads are never joined past the deadline, so a stuck OS call cannot delay
interpreter exit or block the caller (``os.path.isdir`` releases the GIL).
"""
paths = list(paths) # fixed input we can iterate twice; one probe per path
results: dict[str, bool] = {}
def _probe(path: str) -> None:
try:
results[path] = os.path.isdir(path) and os.access(path, os.R_OK)
except OSError:
results[path] = False
threads: list[threading.Thread] = []
for path in paths:
thread = threading.Thread(target = _probe, args = (path,), daemon = True)
thread.start()
threads.append(thread)
deadline = time.monotonic() + timeout
for thread in threads:
thread.join(max(0.0, deadline - time.monotonic()))
# Iterate the fixed input, not results.items(): a probe that timed out is
# still alive and may insert its key here, which would raise "dictionary
# changed size during iteration". results.get() is an atomic read.
return {path for path in paths if results.get(path)}
def _readable_dir_within(path: str, timeout: float) -> bool:
"""``os.path.isdir(path) and os.access(path, R_OK)``, bounded by *timeout* seconds; single-path wrapper over :func:`_readable_dirs_within`."""
return path in _readable_dirs_within((path,), timeout)
def windows_drive_roots(drive_letters: Iterable[str] = string.ascii_uppercase) -> list[Path]:
"""Readable logical drive roots (``C:\\``, ``D:\\`` ...) for the folder browser; the Windows analog of :func:`linux_run_media_mount_roots`.
Without it the allowlist and chips only reach the home drive, so a user
cannot navigate from ``C:`` to ``D:``/``E:``. ``GetLogicalDrives`` drops
unmapped letters; the rest are probed concurrently under a single timeout
and kept only if readable in time. A disconnected mapped drive stays active
in the bitmask and its ``os.path.isdir`` can hang for tens of seconds, so
parallel probing bounds the added delay at ~one timeout rather than one per
drive. Returns ``[]`` off Windows.
"""
if platform.system() != "Windows":
return []
active_mask = _active_windows_drive_bitmask()
candidates: list[str] = []
seen: set[str] = set()
for letter in drive_letters:
letter = letter.strip().rstrip(":").upper()
if len(letter) != 1 or letter not in string.ascii_uppercase:
continue
if active_mask and not active_mask & (1 << (ord(letter) - ord("A"))):
continue
root_text = f"{letter}:\\"
key = os.path.normcase(root_text)
if key in seen:
continue
seen.add(key)
candidates.append(root_text)
# Bounded concurrent probe: an active bitmask bit can still be a
# disconnected mapping whose os.path.isdir blocks, so probe all at once.
readable = _readable_dirs_within(candidates, _DRIVE_PROBE_TIMEOUT_S)
return [Path(root_text) for root_text in candidates if root_text in readable]

View file

@ -0,0 +1,607 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Latest-transformers support check for brand-new model architectures.
When a model's ``model_type`` is absent from every installed transformers overlay
(base 4.57.x plus the .venv_t5_530/550/510 sidecars and, if provisioned, .venv_t5_latest),
Studio cannot load it today. This module answers, without authentication, code execution,
or trust_remote_code:
1. Does the LATEST transformers release on PyPI ship this ``model_type``?
2. Does transformers ``main`` on GitHub ship it (dev-only, not yet installable)?
Sources (all unauthenticated; raw.githubusercontent.com is not API rate-limited and
api.github.com is deliberately never used):
- https://pypi.org/pypi/transformers/json -> latest release version
- https://raw.githubusercontent.com/huggingface/transformers/{ref}/src/transformers/
models/auto/configuration_auto.py + auto_mappings.py -> CONFIG_MAPPING_NAMES
The fetched sources are parsed with the same AST extractor the static router uses
(:func:`utils.transformers_version._model_types_from_source`), so the remote answer is
computed exactly like the local overlay answer.
Results are cached in memory and in a small JSON snapshot under ``studio_root()/cache``
(ttl ~1 day) so repeated tier resolutions never re-fetch; failures are backed off in
memory. Every fetch is bounded (<=5s, one retry), so a hung network cannot block model
loading. Fully offline-safe: offline env vars or the kill switch
``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1`` make every check return None (current
behavior preserved).
The consented install path (:func:`install_latest_transformers`) provisions the
persistent ``.venv_t5_latest`` sidecar via
:func:`utils.transformers_version.ensure_latest_transformers_venv`.
"""
import json
import os
import threading
import time
from pathlib import Path
from loggers import get_logger
from utils.paths.storage_roots import studio_root as _studio_root
from utils.transformers_version import (
_env_offline,
_load_config_json,
_model_types_from_source,
_tier_from_config_mapping,
_config_model_types,
_NESTED_CONFIG_KEYS,
_TIER_RANK,
_model_types_from_config,
_TRANSFORMERS_510_MODEL_TYPES,
_TRANSFORMERS_530_MODEL_TYPES,
_TRANSFORMERS_550_MODEL_TYPES,
ensure_latest_transformers_venv,
latest_venv_pinned_version,
)
logger = get_logger(__name__)
_PYPI_JSON_URL = "https://pypi.org/pypi/transformers/json"
_RAW_URL = (
"https://raw.githubusercontent.com/huggingface/transformers/{ref}"
"/src/transformers/models/auto/{name}"
)
_AUTO_FILES = ("configuration_auto.py", "auto_mappings.py")
_FETCH_TIMEOUT_SECONDS = 5.0
_FETCH_RETRIES = 1
_CACHE_TTL_SECONDS = 24 * 60 * 60
_FAILURE_BACKOFF_SECONDS = 300
_CACHE_FILE_NAME = "transformers_latest_check.json"
_SNAPSHOT_SCHEMA = 1
# Snapshot: {"schema", "fetched_at", "pypi_version", "pypi_model_types", "main_model_types"}.
# Install-in-progress state lives in utils.transformers_version (the sidecar swap reservation).
_lock = threading.Lock()
_memory_snapshot: dict | None = None
_last_failure_at: float = 0.0
_is_fetching: bool = False
_TRUE_VALUES = {"1", "true", "yes", "on"}
def _disabled() -> bool:
"""True if the operator disabled the latest-transformers check entirely."""
return (
os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in _TRUE_VALUES
)
def _cache_file() -> Path:
return _studio_root() / "cache" / _CACHE_FILE_NAME
# Sentinel for HTTP 404 (absent at ref), distinct from transient failures.
_FETCH_MISSING = "__unsloth_fetch_missing__"
def _fetch_text(url: str) -> str | None:
"""GET *url* with a bounded timeout and one retry; None on any failure.
Returns ``_FETCH_MISSING`` (without retrying) on HTTP 404 so callers can tell
"absent at this ref" apart from "network flaked".
"""
import urllib.error
import urllib.request
for attempt in range(1 + _FETCH_RETRIES):
try:
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _FETCH_TIMEOUT_SECONDS) as resp:
return resp.read().decode("utf-8", "replace")
except urllib.error.HTTPError as exc:
if exc.code == 404:
return _FETCH_MISSING
logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc)
except Exception as exc:
logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc)
return None
def _fetch_latest_pypi_version() -> str | None:
"""Latest transformers release version from PyPI's unauthenticated JSON API."""
body = _fetch_text(_PYPI_JSON_URL)
if body is None or body == _FETCH_MISSING:
return None
try:
version = json.loads(body).get("info", {}).get("version")
except Exception as exc:
logger.debug("Could not parse PyPI JSON: %s", exc)
return None
return version if isinstance(version, str) and version else None
def _fetch_remote_model_types(ref: str) -> frozenset[str] | None:
"""CONFIG_MAPPING_NAMES keys at *ref* (a release tag like ``v5.12.0`` or ``main``).
Fetches configuration_auto.py plus auto_mappings.py (the 5.10+ split) from
raw.githubusercontent.com and parses them with the shared AST extractor. A file
that 404s (auto_mappings.py on pre-5.10 tags) is skipped, but a transient fetch
or parse failure of EITHER file fails the whole lookup: most model types live in
auto_mappings.py on current releases, so a partial map cached for the TTL would
make /validate skip the upgrade prompt for architectures the release does ship.
An empty result is likewise a failure so it is never cached as "supports nothing".
"""
keys: set[str] = set()
fetched_any = False
for name in _AUTO_FILES:
source = _fetch_text(_RAW_URL.format(ref = ref, name = name))
if source is None:
return None
if source == _FETCH_MISSING:
continue
fetched_any = True
try:
keys |= _model_types_from_source(source)
except Exception as exc:
logger.debug("Could not parse %s at %s: %s", name, ref, exc)
return None
if not fetched_any or not keys:
return None
return frozenset(keys)
def _load_snapshot_file() -> dict | None:
"""Persisted snapshot from disk, or None (missing/corrupt/old schema)."""
try:
with open(_cache_file(), encoding = "utf-8") as f:
data = json.load(f)
except Exception:
return None
if not isinstance(data, dict) or data.get("schema") != _SNAPSHOT_SCHEMA:
return None
if not isinstance(data.get("fetched_at"), (int, float)):
return None
if not isinstance(data.get("pypi_version"), str):
return None
for key in ("pypi_model_types", "main_model_types"):
value = data.get(key)
if not isinstance(value, list) or not all(isinstance(v, str) for v in value):
return None
return data
def _save_snapshot_file(snapshot: dict) -> None:
"""Atomic best-effort write (tmp + os.replace, Windows-safe); failures only log."""
path = _cache_file()
tmp = path.with_name(path.name + ".tmp")
try:
path.parent.mkdir(parents = True, exist_ok = True)
tmp.write_text(json.dumps(snapshot), encoding = "utf-8")
os.replace(tmp, path)
except Exception as exc:
logger.debug("Could not persist %s: %s", path, exc)
try:
tmp.unlink(missing_ok = True)
except Exception:
pass
def _snapshot_is_fresh(snapshot: dict | None) -> bool:
return (
snapshot is not None
and (time.time() - float(snapshot.get("fetched_at", 0))) < _CACHE_TTL_SECONDS
)
def _refresh_snapshot() -> dict | None:
"""Fetch a fresh snapshot from PyPI + raw.githubusercontent.com; None on failure.
The PyPI version and its tagged mapping are required; the ``main`` mapping is
best-effort (recorded as an empty list plus ``main_checked=False`` when unavailable,
so a dev-only architecture is reported as "unknown" rather than "unsupported").
"""
version = _fetch_latest_pypi_version()
if version is None:
return None
pypi_types = _fetch_remote_model_types(f"v{version}")
if pypi_types is None:
return None
main_types = _fetch_remote_model_types("main")
return {
"schema": _SNAPSHOT_SCHEMA,
"fetched_at": time.time(),
"pypi_version": version,
"pypi_model_types": sorted(pypi_types),
"main_model_types": sorted(main_types) if main_types is not None else [],
"main_checked": main_types is not None,
}
def _get_snapshot() -> dict | None:
"""Current support snapshot: memory -> disk -> network, with TTL and failure backoff.
The network refresh runs outside the lock so a slow fetch cannot stall other
threads in the ASGI pool; _is_fetching deduplicates concurrent refreshes
(losers return None, the graceful fallthrough, rather than waiting).
"""
global _memory_snapshot, _last_failure_at, _is_fetching
with _lock:
if _snapshot_is_fresh(_memory_snapshot):
return _memory_snapshot
disk = _load_snapshot_file()
if _snapshot_is_fresh(disk):
_memory_snapshot = disk
return disk
if _disabled() or _env_offline():
return None
if time.time() - _last_failure_at < _FAILURE_BACKOFF_SECONDS:
return None
if _is_fetching:
return None
_is_fetching = True
fresh = None
try:
fresh = _refresh_snapshot()
finally:
with _lock:
_is_fetching = False
if fresh is None:
_last_failure_at = time.time()
else:
_memory_snapshot = fresh
if fresh is None:
# A stale positive could offer a version PyPI no longer serves; be strict.
return None
_save_snapshot_file(fresh)
return fresh
def clear_caches() -> None:
"""Test helper: drop the in-memory snapshot, failure backoff, and busy flags."""
global _memory_snapshot, _last_failure_at, _is_fetching
with _lock:
_memory_snapshot = None
_last_failure_at = 0.0
_is_fetching = False
from utils.transformers_version import end_sidecar_swap
end_sidecar_swap()
def latest_transformers_supports(model_type: str) -> dict | None:
"""Whether the newest transformers (PyPI release and/or GitHub main) ships *model_type*.
Returns ``{"pypi_version": str, "supported_in_pypi": bool, "supported_in_main": bool}``
or None when the answer is unavailable (offline, kill switch, network failure) the
caller must then fall through to current behavior. Cached (memory + JSON snapshot on
disk, ttl ~1 day) so repeated tier resolutions never re-fetch.
"""
if not isinstance(model_type, str) or not model_type:
return None
if _disabled() or _env_offline():
return None
snapshot = _get_snapshot()
if snapshot is None:
return None
return {
"pypi_version": snapshot["pypi_version"],
"supported_in_pypi": model_type in set(snapshot["pypi_model_types"]),
"supported_in_main": model_type in set(snapshot["main_model_types"]),
}
# model_types the hardcoded tier tables already route; never remote-check these.
def _hardcoded_model_types() -> frozenset[str]:
return frozenset(
_TRANSFORMERS_530_MODEL_TYPES
| _TRANSFORMERS_550_MODEL_TYPES
| _TRANSFORMERS_510_MODEL_TYPES
)
def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dict | None:
"""Upgrade signal for *model_name*, or None when current routing already handles it.
The tier hook for the pre-load ``/validate`` path: fires ONLY when the model's
``model_type`` is absent from every installed overlay (and from the hardcoded tier
tables), i.e. exactly when today's load would fail with an unrecognized-architecture
error. Returns ``{"model_type", "pypi_version", "supported_in_pypi",
"supported_in_main"}`` when the newest transformers knows the type, else None.
Never raises; every network touch is bounded and cached. Offline or with the
``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS`` kill switch it returns None immediately.
"""
try:
if _disabled() or _env_offline():
return None
cfg = _load_config_json(model_name, hf_token)
if not isinstance(cfg, dict):
return None
candidates = _model_types_from_config(cfg)
if not candidates:
return None
# Without a readable base mapping every type looks brand new; bail out.
if not _config_model_types("default"):
return None
hardcoded = _hardcoded_model_types()
missing = [
candidate
for candidate in candidates
if candidate not in hardcoded
and not any(candidate in _config_model_types(tier) for tier in _TIER_RANK)
]
if not missing:
return None
# Latest must load EVERY missing type (wrappers build nested sub-configs
# through CONFIG_MAPPING) or the load still fails.
supports = [latest_transformers_supports(candidate) for candidate in missing]
if any(
s is None or not (s["supported_in_pypi"] or s["supported_in_main"]) for s in supports
):
return None
# Offer the PyPI install only if the release ships every missing type; a
# main-only type in the mix surfaces as dev-only.
model_type = missing[0]
supported_in_pypi = all(s["supported_in_pypi"] for s in supports)
supported_in_main = all(s["supported_in_pypi"] or s["supported_in_main"] for s in supports)
logger.info(
"Model %s has model_type=%s unknown to every installed transformers "
"(latest PyPI %s: %s, main: %s)",
model_name,
model_type,
supports[0]["pypi_version"],
"supported" if supported_in_pypi else "unsupported",
"supported" if supported_in_main else "unsupported",
)
return {
"model_type": model_type,
"pypi_version": supports[0]["pypi_version"],
"supported_in_pypi": supported_in_pypi,
"supported_in_main": supported_in_main,
}
except Exception as exc:
logger.debug("Latest-transformers check failed for '%s': %s", model_name, exc)
return None
# --- Dependency compatibility preflight ------------------------------------------------------
# Sidecars install transformers --no-deps atop the base env. Before installing, compare
# requires_dist: unsatisfied shadowable deps become exact --target pins, anything else blocks.
# Safe to shadow inside the sidecar dir (pure wheels, no torch coupling).
_SHADOWABLE_DEPS = frozenset({"tokenizers", "safetensors"})
# Provided by the sidecar recipe; checked against its pin, not the base env.
_SIDECAR_PROVIDED = {"huggingface-hub": "1.8.0", "hf-xet": "1.4.2"}
# CLI-only; never imported at runtime in Studio's workers.
_IGNORED_DEPS = frozenset({"typer"})
def _canonical_dep_name(name: str) -> str:
return name.lower().replace("_", "-")
def _fetch_requires_dist(version: str) -> list[str] | None:
"""Core (marker-free, non-extra) requires_dist of transformers *version* from PyPI."""
body = _fetch_text(f"https://pypi.org/pypi/transformers/{version}/json")
if body is None or body == _FETCH_MISSING:
return None
try:
reqs = json.loads(body).get("info", {}).get("requires_dist")
except Exception:
return None
if not isinstance(reqs, list):
return None
return [r for r in reqs if isinstance(r, str)]
def _resolve_exact_version(name: str, specifier) -> str | None:
"""Newest PyPI release of *name* satisfying *specifier* (exact pin for the shadow)."""
body = _fetch_text(f"https://pypi.org/pypi/{name}/json")
if body is None or body == _FETCH_MISSING:
return None
try:
from packaging.version import InvalidVersion, Version
releases = json.loads(body).get("releases", {})
best = None
for candidate in releases:
try:
parsed = Version(candidate)
except InvalidVersion:
continue
if parsed.is_prerelease or not specifier.contains(candidate):
continue
if best is None or parsed > Version(best):
best = candidate
return best
except Exception as exc:
logger.debug("Could not resolve an exact %s version: %s", name, exc)
return None
def compat_plan(version: str) -> tuple[tuple[str, ...], list[str]]:
"""(extra exact pins to shadow-install, blocking requirement strings) for *version*.
Compares the release's core requires_dist against the running base env (the env the
workers overlay the sidecar onto). A requirement the base env satisfies needs nothing;
an unsatisfied shadowable dep becomes an exact pin inside the sidecar; any other
unsatisfied requirement is a blocker. An unavailable requires_dist BLOCKS the
install: proceeding unverified could pin a sidecar whose imports then crash the
workers, and the caller just reached PyPI for the version check so a retry is cheap.
"""
reqs = _fetch_requires_dist(version)
if reqs is None:
return (), ["dependency metadata for this release (could not be fetched from PyPI; retry)"]
try:
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as _installed_version
from packaging.requirements import InvalidRequirement, Requirement
except Exception:
return (), []
extras: list[str] = []
blockers: list[str] = []
for raw in reqs:
try:
req = Requirement(raw)
except InvalidRequirement:
continue
if req.extras or (req.marker is not None and not req.marker.evaluate()):
continue
name = _canonical_dep_name(req.name)
if name in _IGNORED_DEPS:
continue
if name in _SIDECAR_PROVIDED:
if not req.specifier.contains(_SIDECAR_PROVIDED[name], prereleases = True):
blockers.append(raw)
continue
try:
installed = _installed_version(req.name)
except PackageNotFoundError:
installed = None
if installed is not None and req.specifier.contains(installed, prereleases = True):
continue
if name in _SHADOWABLE_DEPS:
exact = _resolve_exact_version(name, req.specifier)
if exact is None:
blockers.append(raw)
else:
extras.append(f"{name}=={exact}")
else:
blockers.append(raw)
return tuple(extras), blockers
def is_install_in_progress() -> bool:
"""True while a latest-transformers install or lazy repair holds the sidecar swap
reservation. Training and export starts check this so a fresh worker never
activates the sidecar mid-swap."""
from utils.transformers_version import sidecar_swap_in_progress
return sidecar_swap_in_progress()
def install_latest_transformers(
version: str,
before_swap = None,
reserved: bool = False,
) -> dict:
"""Consented install of the latest transformers sidecar; returns a structured result.
Guards: the requested *version* must match the current PyPI latest from the (cached)
snapshot, so a client cannot pin an arbitrary package version through this endpoint.
On success ``.venv_t5_latest`` is provisioned and pinned; routing then resolves the
new tier automatically on this and every future start. *before_swap* is forwarded
to the stage-and-swap: it runs only after the staged install succeeded, right
before the live sidecar is replaced. *reserved* means the caller already holds the
sidecar swap reservation (the install route takes it before waiting on the
inference lifecycle gate, so worker starts see it for the whole window).
"""
from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap
if not reserved and not try_begin_sidecar_swap():
return {
"success": False,
"version": version,
"message": "A transformers installation is already in progress.",
}
try:
return _install_latest_transformers_locked(version, before_swap = before_swap)
finally:
if not reserved:
end_sidecar_swap()
def _install_latest_transformers_locked(version: str, before_swap = None) -> dict:
"""Body of install_latest_transformers; runs with the in-progress flag held."""
if _disabled():
return {
"success": False,
"version": version,
"message": "Latest-transformers installs are disabled "
"(UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS).",
}
if _env_offline():
return {
"success": False,
"version": version,
"message": "Cannot install: Studio is in offline mode.",
}
# Re-verify against a LIVE snapshot (a release may land inside the cache TTL);
# fall back to the cached one on fetch failure.
global _memory_snapshot
snapshot = _refresh_snapshot()
if snapshot is not None:
with _lock:
_memory_snapshot = snapshot
_save_snapshot_file(snapshot)
else:
snapshot = _get_snapshot()
if snapshot is None:
return {
"success": False,
"version": version,
"message": "Could not verify the latest transformers release on PyPI.",
}
if version != snapshot["pypi_version"]:
return {
"success": False,
"version": version,
"message": f"Requested version {version!r} is not the latest transformers "
f"release ({snapshot['pypi_version']}).",
# Lets the consent dialog retry with the release that superseded the
# one /validate saw, instead of re-sending the stale version forever.
"latest_version": snapshot["pypi_version"],
}
extra_packages, blockers = compat_plan(version)
if blockers:
return {
"success": False,
"version": version,
"message": "Cannot install transformers "
f"{version}: this environment does not satisfy {', '.join(blockers)}. "
"A Studio update is required first.",
}
if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap):
return {
"success": False,
"version": version,
"message": f"Installing transformers {version} failed; see the Studio logs.",
}
_invalidate_capability_caches()
return {
"success": True,
"version": version,
"message": f"Installed transformers {version} into the latest sidecar "
f"(pinned: {latest_venv_pinned_version()}).",
}
def _invalidate_capability_caches():
"""Drop caches computed before the new sidecar existed: tier probes and the
latest tier's model_type mapping (stale on upgrade) plus vision detection
(a raw-heuristic False may now defer to the sidecar AutoConfig probe)."""
try:
from utils import transformers_version as tv
tv._probe_tier_cache.clear()
tv._config_mapping_cache.pop("latest", None)
except Exception:
pass
try:
from utils.models import model_config as mc
mc._vision_detection_cache.clear()
except Exception:
pass

View file

@ -35,9 +35,12 @@ import json
import structlog
from loggers import get_logger
import os
import re
import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path
from utils.native_path_leases import child_env_without_native_path_secret
@ -235,8 +238,12 @@ _VENV_T5_DIR = _VENV_T5_550_DIR
# reuses the workspace torch (torch-agnostic).
_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor")
# Tier precedence: higher rank wins in _higher_tier.
_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3}
# User-consented "latest transformers" sidecar (utils/transformers_latest.py); pinned version in a marker file.
_VENV_T5_LATEST_DIR = str(_studio_root() / ".venv_t5_latest")
_LATEST_PIN_MARKER = ".unsloth_pinned_transformers"
# Tier precedence: higher rank wins in _higher_tier. "latest" outranks every fixed tier.
_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3, "latest": 4}
def _higher_tier(a: str, b: str) -> str:
@ -254,20 +261,40 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None =
``hf_token`` is forwarded to tier detection so a gated/private model whose only 5.x
signal is an authenticated config/tokenizer reaches the right sidecar, not the default.
"""
# Pre-resolve only LoRA adapters; full checkpoints go to get_transformers_tier so their
# local config.json drives the tier (a full checkpoint with a private/offline
# _name_or_path must not resolve to an unreachable HF id and skip its own config).
# Pre-resolve LoRA adapters (local dir or remote adapter repo); full checkpoints
# go to get_transformers_tier so their local config.json drives the tier (a full
# checkpoint with a private/offline _name_or_path must not resolve to an
# unreachable HF id and skip its own config). Remote adapters activate for their
# BASE model, matching latest_tier_active_for and the inference worker.
if _is_lora_adapter_dir(Path(model_name)):
resolved = _resolve_base_model(model_name)
else:
resolved = model_name
resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name
tier = get_transformers_tier(resolved, hf_token)
if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"):
# Gate on a real local config.json: a checkpoint carries config the base may not
# surface, but path names alone must not upgrade a plain adapter.
tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token))
if tier == "510":
if tier == "latest":
pinned = latest_venv_pinned_version()
if pinned is None or not _ensure_venv_t5_latest_exists():
raise RuntimeError(
f"Cannot activate the latest-transformers sidecar: "
f".venv_t5_latest missing or unpinned at {_VENV_T5_LATEST_DIR}"
)
if _VENV_T5_LATEST_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_LATEST_DIR)
logger.info(
"Prepended transformers %s venv to sys.path from %s "
"(path only; the loaded version is confirmed later by "
"'Subprocess loaded transformers ...' on first import)",
pinned,
_VENV_T5_LATEST_DIR,
)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_LATEST_DIR + (os.pathsep + _pp if _pp else "")
elif tier == "510":
if not _ensure_venv_t5_510_exists():
raise RuntimeError(
f"Cannot activate transformers {TRANSFORMERS_510_VERSION}: "
@ -322,6 +349,34 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None =
logger.info("Using default transformers (4.57.x) for %s", model_name)
def latest_tier_active_for(model_name: str, hf_token: str | None = None) -> bool:
"""True when *model_name* routes to the consented latest-transformers sidecar.
Mirrors the inference worker's pre-activation resolution (local adapter dir,
then a remote adapter's Hub adapter_config.json). ``latest`` only wins when
the sidecar exists with a valid pin, i.e. exactly the loads that will import
the newest release. Never raises: any resolution failure returns False so
callers treat the model as a known tier.
"""
try:
# No consented sidecar pin means nothing routes to latest; return before
# any resolution so the common case costs no config or network reads.
if latest_venv_pinned_version() is None:
return False
if _is_lora_adapter_dir(Path(model_name)):
resolved = _resolve_base_model(model_name)
else:
# A remote LoRA activates the sidecar for its BASE model; sizing and the
# worker's 4-bit guard must see that base too, not the adapter repo.
resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name
tier = get_transformers_tier(resolved, hf_token)
if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"):
tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token))
return tier == "latest"
except Exception:
return False
def _has_adapter_weights(path: Path) -> bool:
"""True if *path* holds LoRA adapter weight files (``adapter_model.*``)."""
try:
@ -881,17 +936,85 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None:
_config_mapping_cache: dict[str, frozenset[str]] = {}
def _latest_tier_disabled() -> bool:
"""Kill switch shared with utils.transformers_latest: lets operators roll
back a provisioned latest sidecar without deleting files."""
return os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in (
"1",
"true",
"yes",
"on",
)
# Failed lazy repairs back off so a broken sidecar can't turn every routing
# call into a pip install attempt.
_latest_repair_failed_at: float = 0.0
_LATEST_REPAIR_BACKOFF_SECS = 5 * 60
def _latest_sidecar_intact() -> bool:
"""The pinned latest sidecar exists with its transformers dir and every pinned
package. False when the pin itself is gone: a cached 'latest' mapping must then be
dropped (routing re-resolves to no latest tier), not trusted, and a sidecar that kept
transformers/ but lost a pinned package must self-heal rather than route models to a
latest tier that fails activation in workers, which refuse parent-only repairs.
(_overlay_transformers_dir only calls this after gating on a present pin, so the
pin-missing case here is the cache-revalidation caller whose pin was deleted after
the mapping was first cached.)"""
pin = _latest_pin_data()
if pin is None:
return False
return _venv_dir_is_valid(_VENV_T5_LATEST_DIR, tuple(pin["packages"]))
def _overlay_transformers_dir(tier: str) -> str | None:
"""transformers source dir for a tier, located without importing it."""
global _latest_repair_failed_at
if tier != "default":
root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier)
# latest requires a valid pin and the kill switch off.
if tier == "latest" and (_latest_tier_disabled() or latest_venv_pinned_version() is None):
return None
root = {
"530": _VENV_T5_530_DIR,
"550": _VENV_T5_550_DIR,
"510": _VENV_T5_510_DIR,
"latest": _VENV_T5_LATEST_DIR,
}.get(tier)
src = os.path.join(root, "transformers") if root else None
if src and tier == "latest" and not _latest_sidecar_intact():
# A valid pin whose sidecar vanished or lost a pinned package (partial
# deletion, disk issue, interrupted external edits) must self-heal, or
# latest-only models either silently route to older tiers or reach a
# worker that cannot repair, failing every load until a manual
# reinstall. Repair under the swap reservation; back off after a
# failure so routing calls don't hammer pip.
repaired = False
if time.time() - _latest_repair_failed_at >= _LATEST_REPAIR_BACKOFF_SECS:
if _ensure_venv_t5_latest_exists():
_latest_repair_failed_at = 0.0
repaired = True
else:
_latest_repair_failed_at = time.time()
if not repaired:
# Still broken: treat the overlay as unavailable rather than route
# models to a tier whose worker activation is known to fail. Models
# an older tier supports keep loading there until a repair succeeds,
# matching the behavior when the sidecar dir is missing entirely.
return None
return src if src and _safe_is_dir(Path(src)) else None
# default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one
# is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir.
sidecars = tuple(
os.path.abspath(d) + os.sep
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR)
for d in (
_VENV_T5_530_DIR,
_VENV_T5_550_DIR,
_VENV_T5_510_DIR,
_VENV_T5_LATEST_DIR,
_VENV_LLMCOMPRESSOR_DIR,
)
)
candidates = []
try:
@ -930,11 +1053,47 @@ def _mapping_first_keys(value: ast.AST) -> set[str]:
return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)}
def _model_types_from_source(source: str) -> set[str]:
"""model_type keys of CONFIG_MAPPING_NAMES in *source* (AST only, no execution).
Handles the direct ``CONFIG_MAPPING_NAMES = ...`` binding (dict literal or
OrderedDict/dict call over 2-tuple lists and **{...} unpacking) and any
``CONFIG_MAPPING_NAMES.update({...})`` mutation. Shared by the on-disk overlay
reader below and the remote latest-release checker (utils/transformers_latest.py).
"""
keys: set[str] = set()
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets
):
keys |= _mapping_first_keys(node.value)
elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
fn = node.value.func
if (
isinstance(fn, ast.Attribute)
and fn.attr == "update"
and isinstance(fn.value, ast.Name)
and fn.value.id == "CONFIG_MAPPING_NAMES"
):
keys |= _mapping_first_keys(node.value)
return keys
def _config_model_types(tier: str) -> frozenset[str]:
"""model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py)."""
# Kill switch beats the cache: a stale mapping must not keep routing latest-only models until restart.
if tier == "latest" and _latest_tier_disabled():
return frozenset()
cached = _config_mapping_cache.get(tier)
if cached is not None:
return cached
# A cached 'latest' mapping can outlive the sidecar it was parsed from: if the
# pinned sidecar was since deleted or lost a package in this process, drop the
# cache so routing re-resolves through _overlay_transformers_dir (which self-heals)
# instead of routing latest-only models to a broken tier until restart.
if tier != "latest" or _latest_sidecar_intact():
return cached
_config_mapping_cache.pop("latest", None)
tdir = _overlay_transformers_dir(tier)
if tdir is None:
return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads
@ -944,22 +1103,7 @@ def _config_model_types(tier: str) -> frozenset[str]:
if not _safe_is_file(path):
continue
try:
tree = ast.parse(path.read_text(encoding = "utf-8"))
for node in ast.walk(tree):
# direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation
if isinstance(node, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets
):
keys |= _mapping_first_keys(node.value)
elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
fn = node.value.func
if (
isinstance(fn, ast.Attribute)
and fn.attr == "update"
and isinstance(fn.value, ast.Name)
and fn.value.id == "CONFIG_MAPPING_NAMES"
):
keys |= _mapping_first_keys(node.value)
keys |= _model_types_from_source(path.read_text(encoding = "utf-8"))
except Exception:
continue
result = frozenset(keys)
@ -967,23 +1111,73 @@ def _config_model_types(tier: str) -> frozenset[str]:
return result
def _tier_from_config_mapping(cfg: dict) -> str | None:
"""Lowest tier whose transformers ships cfg's model_type, or None if unknown."""
model_type = cfg.get("model_type")
if not isinstance(model_type, str):
for key in _NESTED_CONFIG_KEYS:
sub = cfg.get(key)
if isinstance(sub, dict) and isinstance(sub.get("model_type"), str):
model_type = sub["model_type"]
break
if not isinstance(model_type, str):
return None
def _model_types_from_config(cfg: dict) -> list[str]:
"""All model_types in the config: the primary (top-level, else first nested)
first, then every other nested sub-config. Wrappers instantiate sub-configs
through CONFIG_MAPPING, so nested types matter for routing too."""
seen: list[str] = []
def add(value):
if isinstance(value, str) and value and value not in seen:
seen.append(value)
add(cfg.get("model_type"))
for key in _NESTED_CONFIG_KEYS:
sub = cfg.get(key)
if isinstance(sub, dict):
add(sub.get("model_type"))
for value in cfg.values():
if isinstance(value, dict):
add(value.get("model_type"))
return seen
def _lowest_tier_for(model_type: str) -> str | None:
for tier in sorted(_TIER_RANK, key = _TIER_RANK.get):
if model_type in _config_model_types(tier):
return tier
return None
def _tier_from_config_mapping(cfg: dict) -> str | None:
"""Lowest tier able to load every model_type in cfg, or None when the
primary type is unknown everywhere. A nested type can raise the tier (its
sub-config is built through CONFIG_MAPPING); an unknown nested type never
vetoes, since no installed tier could load it either way (the latest
checker handles surfacing the install prompt for it)."""
types = _model_types_from_config(cfg)
if not types:
return None
best = _lowest_tier_for(types[0])
if best is None:
return None
for model_type in types[1:]:
tier = _lowest_tier_for(model_type)
if tier is not None and _TIER_RANK[tier] > _TIER_RANK[best]:
best = tier
return best
def _raise_tier_for_nested(cfg: dict | None, tier: str) -> str:
"""Raise *tier* when the mapping resolver needs a higher one for *cfg*.
A wrapper's top-level model_type can match a hardcoded fast path while a
nested text/vision config's type only exists in a newer sidecar (e.g. the
installed latest); its sub-config is built through CONFIG_MAPPING, so the
fast-path tier would fail to load it. Raise-only: never lowers a fast-path
match, so name overrides (Qwen3.6) keep their tier. Never raises an
exception: a resolution failure keeps the fast-path tier."""
if not isinstance(cfg, dict):
return tier
try:
mapped = _tier_from_config_mapping(cfg)
if mapped is not None and _TIER_RANK.get(mapped, 0) > _TIER_RANK.get(tier, 0):
return mapped
except Exception:
pass
return tier
# --- AutoConfig probe: general tier resolution for ambiguous models ----------
# When the cheap signals only say "needs some 5.x", parse config.json with the built-in
# parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond
@ -1039,9 +1233,19 @@ def _probe_tier_venvs():
"530": (_VENV_T5_530_DIR, _ensure_venv_t5_530_exists),
"550": (_VENV_T5_550_DIR, _ensure_venv_t5_550_exists),
"510": (_VENV_T5_510_DIR, _ensure_venv_t5_510_exists),
"latest": (_VENV_T5_LATEST_DIR, _ensure_venv_t5_latest_exists),
}
def _probe_tier_order() -> tuple[str, ...]:
"""Sidecar probe order. The consented "latest" sidecar joins only once it is
provisioned (pin marker present): an absent optional tier must not flip the probe's
skipped-tier bookkeeping, keeping pre-latest behavior byte-identical."""
if not _latest_tier_disabled() and latest_venv_pinned_version() is not None:
return _PROBE_TIER_ORDER + ("latest",)
return _PROBE_TIER_ORDER
def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) -> bool | None:
"""Parse config.json with the built-in parser inside *target_dir*'s sidecar.
True = parses, False = parse/version failure (escalate), None = transient
@ -1119,7 +1323,7 @@ def _probe_tier(
stays on the default. Cached per _probe_cache_key (process lifetime). No Hub sha is
resolved: that would import huggingface_hub before the sidecar is on sys.path.
"""
if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes"):
if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes", "on"):
return floor
key = _probe_cache_key(model_name)
# Key by probe mode: the default-first path can return 'default', which must not be
@ -1127,7 +1331,10 @@ def _probe_tier(
if include_default or floor != "530":
key = f"{key}\0floor={floor}:def={int(include_default)}"
if key in _probe_tier_cache:
return _probe_tier_cache[key]
cached = _probe_tier_cache[key]
# Kill switch beats the cache (like _config_model_types): a stale 'latest' probe must not keep activating it.
if cached != "latest" or not _latest_tier_disabled():
return cached
def _cache(tier: str, *, skipped: bool) -> str:
# Do not pin a result that depended on a skipped lower tier: once that sidecar is
@ -1137,7 +1344,8 @@ def _probe_tier(
return tier
venvs = _probe_tier_venvs()
order = (("default",) + _PROBE_TIER_ORDER) if include_default else _PROBE_TIER_ORDER
sidecar_order = _probe_tier_order()
order = (("default",) + sidecar_order) if include_default else sidecar_order
probed_count = 0
skipped_any = False
for tier in order:
@ -1264,17 +1472,21 @@ def get_transformers_tier(
cfg = _load_config_json(model_name, hf_token)
if cfg is not None:
if _config_needs_510(cfg):
tier = _raise_tier_for_nested(cfg, "510")
logger.info(
"Transformers tier 510 selected for %s (local config.json check)",
"Transformers tier %s selected for %s (local config.json check)",
tier,
model_name,
)
return "510"
return tier
if _config_needs_550(cfg):
tier = _raise_tier_for_nested(cfg, "550")
logger.info(
"Transformers tier 550 selected for %s (local config.json check)",
"Transformers tier %s selected for %s (local config.json check)",
tier,
model_name,
)
return "550"
return tier
if _config_needs_530(cfg):
# Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name. Only a real
# Hub id (or the folder basename) may override 530, so a stale local
@ -1287,17 +1499,20 @@ def get_transformers_tier(
)
override = _higher_tier_name_override(hint_src)
if override is not None:
override = _raise_tier_for_nested(cfg, override)
logger.info(
"Transformers tier %s selected for %s (name overrides 530 config)",
override,
model_name,
)
return override
tier = _raise_tier_for_nested(cfg, "530")
logger.info(
"Transformers tier 530 selected for %s (local config.json check)",
"Transformers tier %s selected for %s (local config.json check)",
tier,
model_name,
)
return "530"
return tier
# Unknown arch: resolve the base id from config. A resolved local dir
# recurses (config check); a Hub id uses name rules only (no network).
resolved = _resolve_base_model(model_name)
@ -1359,6 +1574,13 @@ def get_transformers_tier(
result = _tier_from_name(model_name)
if result is not None:
tier, match = result
# With a consented latest sidecar pinned, a name that matches a fixed
# tier can still carry a latest-only model_type (e.g. a newer variant
# reusing a family name); consult the config so an accepted upgrade
# actually routes to the sidecar it installed. Costs a config read only
# in the pinned case, keeping the pre-latest path I/O-free.
if latest_venv_pinned_version() is not None:
tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), tier)
logger.info(
"Transformers tier %s selected for %s (substring match: %s)",
tier,
@ -1369,11 +1591,13 @@ def get_transformers_tier(
# --- Slow config fallbacks (network for HF IDs; authenticated with hf_token) --------
if _check_config_needs_510(model_name, hf_token):
logger.info("Transformers tier 510 selected for %s (config.json check)", model_name)
return "510"
tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "510")
logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
return tier
if _check_config_needs_550(model_name, hf_token):
logger.info("Transformers tier 550 selected for %s (config.json check)", model_name)
return "550"
tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "550")
logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
return tier
if _check_config_needs_530(model_name, hf_token):
# Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name; honor a real Hub-id name
# hint from _name_or_path before selecting 530.
@ -1383,14 +1607,16 @@ def get_transformers_tier(
base if isinstance(base, str) and base != model_name else None
)
if override is not None:
override = _raise_tier_for_nested(remote_cfg, override)
logger.info(
"Transformers tier %s selected for %s (name overrides 530 config)",
override,
model_name,
)
return override
logger.info("Transformers tier 530 selected for %s (config.json check)", model_name)
return "530"
tier = _raise_tier_for_nested(remote_cfg, "530")
logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
return tier
# _load_config_json (not the cache-only reader) so a config served from the hub
# cache during a transient outage still feeds the mapping resolver.
remote_cfg = _load_config_json(model_name, hf_token)
@ -1657,6 +1883,471 @@ def _ensure_venv_t5_exists() -> bool:
return _ensure_venv_t5_550_exists()
# --- User-consented "latest transformers" sidecar (.venv_t5_latest) --------------------------
# Provisioned via ensure_latest_transformers_venv() after the user confirms the upgrade popup
# (utils/transformers_latest.py); pinned in a marker file so restarts revalidate and routing auto-picks it.
# PEP 440-ish release strings only (guards the pip install spec against injection).
_LATEST_VERSION_RE = r"[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?"
def _is_valid_version_string(version: str) -> bool:
import re
return isinstance(version, str) and re.fullmatch(_LATEST_VERSION_RE, version) is not None
# Only the sidecar recipe's own packages, as plain (optionally ==pinned) specs, may
# come from the on-disk pin marker; anything else (URLs, extras, options) is rebuilt.
_PIN_SPEC_RE = re.compile(r"^[A-Za-z0-9_.-]+(==[A-Za-z0-9_.+-]+)?$")
_PIN_ALLOWED_NAMES = frozenset(
{
"transformers",
"huggingface_hub",
"huggingface-hub",
"hf_xet",
"hf-xet",
"tiktoken",
"tokenizers",
"safetensors",
}
)
def _is_safe_pin_spec(spec: str) -> bool:
if not _PIN_SPEC_RE.match(spec):
return False
name = spec.split("==", 1)[0].lower().replace("_", "-")
return name in {n.replace("_", "-") for n in _PIN_ALLOWED_NAMES}
def _recover_stranded_latest_sidecar() -> None:
"""Restore a sidecar stranded at ``.old`` by a swap whose activation rename AND its
rollback both failed (e.g. a lingering worker file handle on Windows blocked both).
That double failure leaves no live dir and the pin marker gone with it, so the
sidecar reads as unprovisioned and never self-heals. Recover only when no live dir
exists and no swap is in flight: the reservation is held throughout the swap, so the
transient live-absent window of a legitimate swap never triggers a restore."""
live = Path(_VENV_T5_LATEST_DIR)
retired = Path(_VENV_T5_LATEST_DIR + ".old")
try:
if live.exists() or not retired.is_dir() or sidecar_swap_in_progress():
return
os.rename(retired, live)
logger.info("Recovered .venv_t5_latest from a stranded .old after a failed swap")
except OSError:
pass
def _latest_pin_data() -> dict | None:
"""Parsed pin marker: {"version": str, "packages": [specs...]}, or None.
The marker is JSON; a plain version string (older/simpler writers) is tolerated and
expanded with the default package set.
"""
_recover_stranded_latest_sidecar()
marker = Path(_VENV_T5_LATEST_DIR) / _LATEST_PIN_MARKER
try:
if not marker.is_file():
return None
raw = marker.read_text(encoding = "utf-8").strip()
except Exception:
return None
try:
data = json.loads(raw)
except ValueError:
data = raw
if isinstance(data, str):
if not _is_valid_version_string(data):
return None
return {"version": data, "packages": list(_venv_t5_latest_packages(data))}
if not isinstance(data, dict):
return None
version = data.get("version")
if not _is_valid_version_string(version):
return None
packages = data.get("packages")
if not (
isinstance(packages, list)
and packages
and all(isinstance(p, str) and _is_safe_pin_spec(p) for p in packages)
):
# Malformed or unexpected specs (the pin is user-writable on disk) never
# reach pip: rebuild the canonical set for the pinned version instead.
packages = list(_venv_t5_latest_packages(version))
return {"version": version, "packages": packages}
def latest_venv_pinned_version() -> str | None:
"""Exact transformers version pinned in .venv_t5_latest's marker, or None if the
sidecar was never provisioned (or the marker is unreadable/invalid)."""
data = _latest_pin_data()
return data["version"] if data else None
def _venv_t5_latest_packages(version: str, extra_packages: tuple[str, ...] = ()) -> tuple[str, ...]:
"""Package set for the latest sidecar; mirrors the fixed .venv_t5_* sidecars.
*extra_packages* carries dep-compat shadows (e.g. a newer tokenizers) computed by
utils.transformers_latest before install."""
return (
f"transformers=={version}",
"huggingface_hub==1.8.0",
"hf_xet==1.4.2",
"tiktoken",
) + tuple(extra_packages)
# Single reservation for ANY .venv_t5_latest replacement (consented install or lazy repair),
# checked by training/export starts so no worker spawns mid-swap. Backed by a lock FILE (not just
# this flag) so a lazy repair running in a worker subprocess stays visible to the parent's route
# checks; the in-process flag marks ownership (only the owner unlinks the file).
_sidecar_swap_lock = threading.Lock()
_sidecar_swap_active = False
_sidecar_swap_token: str | None = None
_sidecar_swap_kind: str | None = None
# An install is minutes; a lock this old is a crashed owner, not a live swap.
_SWAP_LOCK_STALE_SECS = 2 * 60 * 60
def _swap_lock_path() -> Path:
return Path(_VENV_T5_LATEST_DIR + ".swaplock")
def _pid_alive(pid) -> bool:
if not isinstance(pid, int) or pid <= 0:
return False
try:
import psutil
return psutil.pid_exists(pid)
except Exception:
pass
if os.name == "nt":
# os.kill(pid, 0) is NOT a POSIX signal-0 liveness probe on Windows: signal 0
# is CTRL_C_EVENT, so CPython routes it through GenerateConsoleCtrlEvent (a real
# Ctrl+C to that console group) rather than a harmless check. Probe via OpenProcess.
try:
import ctypes
from ctypes import wintypes
kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
kernel32.OpenProcess.restype = wintypes.HANDLE
# PROCESS_QUERY_LIMITED_INFORMATION: minimal right, granted across integrity levels.
handle = kernel32.OpenProcess(0x1000, False, pid)
if handle:
kernel32.CloseHandle(handle)
return True
# ERROR_ACCESS_DENIED means the process exists but we may not query it.
return ctypes.get_last_error() == 5
except Exception:
return False
try:
os.kill(pid, 0)
return True
except OSError:
return False
except Exception:
return False
def _swap_lock_is_stale(path: Path) -> bool:
"""Stale when the recorded owner is provably dead: a crashed installer is reclaimed
at once, not after the long cutoff, so `/load`, training, export, and repair are not
wedged for hours after a crash. A live but slow pip install keeps its lock (its PID
is alive), so breaking it and racing two swaps on the same staging dirs stays
impossible. Only a lock whose PID can't be read (mid-write or corrupt) falls back to
the age cutoff, so the create-before-metadata-write window is never mistaken for dead."""
try:
age = time.time() - path.stat().st_mtime
except OSError:
return False
data = _read_swap_lock(path) or {}
pid = data.get("pid")
if not isinstance(pid, int) or pid <= 0:
return age > _SWAP_LOCK_STALE_SECS
return not _pid_alive(pid)
class SidecarSwapInProgress(RuntimeError):
"""A worker start lost the race to a .venv_t5_latest install/repair; retryable."""
def _read_swap_lock(path: Path) -> dict | None:
try:
data = json.loads(path.read_text(encoding = "utf-8"))
return data if isinstance(data, dict) else {}
except FileNotFoundError:
return None
except OSError:
return {}
except Exception:
return {}
def try_begin_sidecar_swap(kind: str = "install") -> bool:
"""Reserve the sidecar swap window; False when one is already reserved
(in this process or, via the lock file, in any worker subprocess).
*kind* is "install" (consented route) or "repair" (lazy venv repair)."""
global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind
with _sidecar_swap_lock:
if _sidecar_swap_active:
return False
token = f"{os.getpid()}-{time.time_ns()}"
path = _swap_lock_path()
try:
path.parent.mkdir(parents = True, exist_ok = True)
except OSError:
pass
for attempt in range(2):
try:
fd = os.open(str(path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
break
except FileExistsError:
if attempt or not _swap_lock_is_stale(path):
return False
try:
path.unlink()
except OSError:
return False
except OSError:
# Lock file not creatable (odd filesystem): fall back to the process-local reservation.
fd = None
break
if fd is not None:
try:
with os.fdopen(fd, "w") as f:
f.write(
json.dumps(
{"pid": os.getpid(), "at": time.time(), "token": token, "kind": kind}
)
)
except OSError:
pass
_sidecar_swap_active = True
_sidecar_swap_token = token
_sidecar_swap_kind = kind
return True
def end_sidecar_swap() -> None:
"""Release the reservation taken by :func:`try_begin_sidecar_swap`."""
global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind
with _sidecar_swap_lock:
if _sidecar_swap_active:
# Only the file WE wrote is removed: if this reservation was declared
# stale and superseded, unlinking blindly would drop the new owner's
# live lock and unguard its in-flight swap.
path = _swap_lock_path()
data = _read_swap_lock(path)
if data is not None and data.get("token", _sidecar_swap_token) == _sidecar_swap_token:
try:
path.unlink()
except OSError:
pass
_sidecar_swap_active = False
_sidecar_swap_token = None
_sidecar_swap_kind = None
def sidecar_swap_in_progress() -> bool:
"""True while a .venv_t5_latest install or repair holds the reservation,
in this process or any other Studio process (lock file)."""
return sidecar_swap_kind() is not None
def sidecar_swap_kind() -> str | None:
"""The active reservation's kind ("install" / "repair"), or None when idle.
Lets guards that rely on the install route's own abort-on-active-worker
checks keep refusing for repairs, which have no such checks."""
with _sidecar_swap_lock:
if _sidecar_swap_active:
return _sidecar_swap_kind or "install"
path = _swap_lock_path()
try:
if not path.is_file() or _swap_lock_is_stale(path):
return None
except OSError:
return None
data = _read_swap_lock(path) or {}
kind = data.get("kind")
return kind if kind in ("install", "repair") else "install"
def _stage_and_swap_latest_venv(
version: str,
packages: tuple[str, ...],
before_swap = None,
) -> bool:
"""Stage-and-swap: build the new sidecar next to the live one and swap only
once complete, so a failed install or marker write never destroys a
previously working .venv_t5_latest or its pin. Shared by the consented
install and the lazy repair path. *before_swap* (optional callable) runs
after the staging build succeeds and immediately before the live dir is
replaced, so callers can tear down workers only when the swap is certain;
if it raises, the previous sidecar is left untouched."""
staging = _VENV_T5_LATEST_DIR + ".staging"
retired = _VENV_T5_LATEST_DIR + ".old"
shutil.rmtree(staging, ignore_errors = True)
try:
if not _ensure_venv_dir(staging, packages, f"transformers {version} (latest)"):
# No exception, so the except cleanup below never runs; drop the partial dir.
shutil.rmtree(staging, ignore_errors = True)
return False
(Path(staging) / _LATEST_PIN_MARKER).write_text(
json.dumps({"version": version, "packages": list(packages)}), encoding = "utf-8"
)
if before_swap is not None:
before_swap()
shutil.rmtree(retired, ignore_errors = True)
if os.path.isdir(_VENV_T5_LATEST_DIR):
os.rename(_VENV_T5_LATEST_DIR, retired)
try:
os.rename(staging, _VENV_T5_LATEST_DIR)
except OSError:
# Restore the previous sidecar if the final swap fails.
if not os.path.isdir(_VENV_T5_LATEST_DIR) and os.path.isdir(retired):
os.rename(retired, _VENV_T5_LATEST_DIR)
raise
except Exception as exc:
logger.error("Could not provision transformers %s into .venv_t5_latest: %s", version, exc)
shutil.rmtree(staging, ignore_errors = True)
return False
shutil.rmtree(retired, ignore_errors = True)
# CONFIG_MAPPING_NAMES may have changed: drop the cached key set.
_config_mapping_cache.pop("latest", None)
logger.info("Provisioned .venv_t5_latest with transformers %s", version)
return True
def _workers_active_for_repair() -> bool:
"""Best-effort: any parent-visible chat/training/export worker alive. Never
raises; unavailable backends (worker subprocess, early startup) count idle."""
try:
from core.training import get_training_backend
if get_training_backend().is_training_active():
return True
except Exception:
pass
try:
from core.export import get_export_backend
_export = get_export_backend()
if _export.is_export_active():
return True
_alive = getattr(_export, "is_worker_alive", None)
if callable(_alive) and _alive():
return True
except Exception:
pass
try:
from core.inference import get_inference_backend
backend = get_inference_backend()
if getattr(backend, "active_model_name", None):
return True
# An in-flight load counts too: its worker spawns moments later.
if getattr(backend, "loading_models", None):
return True
_alive = getattr(backend, "is_worker_alive", None)
if callable(_alive) and _alive():
return True
except Exception:
pass
return False
def _ensure_venv_t5_latest_exists() -> bool:
"""Ensure .venv_t5_latest/ holds its pinned transformers version.
Never installs without a pin: an unprovisioned sidecar (no marker) returns False so
routing and probing behave exactly as before the feature existed. With a pin present
it repairs a broken dir the same way the fixed sidecars do.
"""
pin = _latest_pin_data()
if pin is None:
return False
version = pin["version"]
packages = tuple(pin["packages"])
if _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages):
return True
if _env_offline():
logger.warning(
".venv_t5_latest (transformers %s) is incomplete and offline mode is set; "
"cannot repair it.",
version,
)
return False
# Repairs are a parent-process action: a worker child's backend singletons are
# empty, so it cannot see live siblings that may still lazy-import from the
# sidecar. Fail activation in the child instead; the parent's routing
# self-heal (guarded below) performs the actual repair.
try:
import multiprocessing as _mp
if _mp.parent_process() is not None:
logger.warning(
".venv_t5_latest is incomplete; repairs run in the parent process. "
"Retry after the parent repairs the sidecar."
)
return False
except Exception:
pass
# Same stage-and-swap as the install, under the same reservation so training/export starts
# (which check sidecar_swap_in_progress) wait out a lazy repair; a failed repair keeps the pin.
if not try_begin_sidecar_swap(kind = "repair"):
logger.warning(
"Cannot repair .venv_t5_latest: another sidecar install or repair is in progress."
)
return False
try:
# Worker check UNDER the reservation (the install route quiesces workers;
# a repair has none): worker starts set their active markers BEFORE
# rechecking the reservation, so either this check sees them and aborts,
# or their recheck sees this reservation and aborts -- no interleaving
# lets a worker spawn against a mid-swap sidecar.
if _workers_active_for_repair():
logger.warning(
"Cannot repair .venv_t5_latest: active chat/training/export workers "
"may be importing from it. Retry when they are idle."
)
return False
return _stage_and_swap_latest_venv(version, packages)
finally:
end_sidecar_swap()
def ensure_latest_transformers_venv(
version: str,
extra_packages: tuple[str, ...] = (),
before_swap = None,
) -> bool:
"""Provision .venv_t5_latest/ pinned to *version* (user-consented install path).
Reuses the same --target/--no-deps installer as the fixed sidecars, then writes the pin
marker (version + full package set) so the venv persists across restarts and
:func:`latest_venv_pinned_version` / routing pick it up automatically.
*extra_packages* carries dep-compat shadows (see utils.transformers_latest).
Returns True on success.
"""
if not _is_valid_version_string(version):
logger.error("Refusing to install invalid transformers version %r", version)
return False
if _env_offline():
logger.warning(
"Cannot install transformers %s: HF/transformers offline mode is set.", version
)
return False
packages = _venv_t5_latest_packages(version, extra_packages)
pin = _latest_pin_data()
if (
pin is not None
and pin["version"] == version
and tuple(pin["packages"]) == packages
and _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages)
):
return True
return _stage_and_swap_latest_venv(version, packages, before_swap = before_swap)
# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) ---------------------
# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize
# Qwen3.5 / Gemma-4 / Llama.
@ -1819,7 +2510,7 @@ def _activate_venv(venv_dir: str, label: str) -> None:
def _deactivate_5x() -> None:
"""Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport."""
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR):
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_T5_LATEST_DIR):
while d in sys.path:
sys.path.remove(d)
logger.info("Removed venv_t5 dirs from sys.path")
@ -1853,14 +2544,25 @@ def ensure_transformers_version(model_name: str) -> None:
if _is_lora_adapter_dir(Path(model_name)):
resolved = _resolve_base_model(model_name)
else:
resolved = model_name
# A remote adapter's tier is its BASE model's (see activation above).
resolved = _remote_lora_base(model_name) or model_name
tier = get_transformers_tier(resolved)
if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"):
# Gate on a real local config.json: a checkpoint carries config the base may not
# surface, but path names alone must not upgrade a plain adapter.
tier = _higher_tier(tier, get_transformers_tier(model_name))
if tier == "510":
if tier == "latest":
pinned = latest_venv_pinned_version()
if pinned is None:
raise RuntimeError(
f"Cannot activate the latest-transformers sidecar: "
f"no pin marker at {_VENV_T5_LATEST_DIR}"
)
target_version = pinned
venv_dir = _VENV_T5_LATEST_DIR
ensure_fn = _ensure_venv_t5_latest_exists
elif tier == "510":
target_version = TRANSFORMERS_510_VERSION
venv_dir = _VENV_T5_510_DIR
ensure_fn = _ensure_venv_t5_510_exists

View file

@ -16,6 +16,7 @@ import {
type ChatSearch,
} from "@/features/chat";
import { RemoteCodeConsentDialog } from "@/features/security";
import { TransformersUpgradeDialog } from "@/features/transformers-upgrade";
import { useTrainingUnloadGuard } from "@/features/training";
import { useExportRuntimeLifecycle } from "@/features/export";
import { hasAuthToken } from "@/features/auth";
@ -230,6 +231,7 @@ function RootLayout() {
<PersonalizationSyncMount />
{!isAuthFlowRoute && <SettingsDialog />}
<RemoteCodeConsentDialog />
<TransformersUpgradeDialog />
{hideNavbar ? (
<main className="flex-1 pt-[var(--studio-hidden-route-top-inset,0px)] [--studio-titlebar-height:var(--studio-hidden-route-top-inset,0px)]">
<Suspense fallback={<RouteFallback />}>

View file

@ -79,6 +79,7 @@ import { McpComposerButton } from "@/features/chat/mcp-composer-button";
import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities";
import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled";
import { BypassPermissionsMenuItem } from "@/features/chat/bypass-permissions-menu-item";
import { PermissionModeComposerPill } from "@/features/chat/permission-mode-select";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store";
import { PROMPT_QUEUE_STOP_EVENT } from "@/features/chat/utils/prompt-queue-boundary";
@ -131,7 +132,6 @@ import {
Image03Icon,
McpServerIcon,
PencilRulerIcon,
ShieldBanIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useNavigate } from "@tanstack/react-router";
@ -1428,11 +1428,14 @@ const Composer: FC<{
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
// More than 4 pills: collapse to icons only. Search and Code always show;
// More than 4 pills: collapse to icons only. Search and Code always show; the
// permission pill shows in every mode except "off" (it renders null there);
// Images, RAG, Canvas and MCP are conditional.
const pillsCompact =
2 +
(permissionMode !== "off" ? 1 : 0) +
(ragEnabled ? 1 : 0) +
(supportsBuiltinImageGeneration ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
@ -1856,9 +1859,9 @@ const Composer: FC<{
data-pill-compact={pillsCompact ? "true" : undefined}
>
<ComposerToolsMenu side={effectiveMenuSide} />
{/* Active-mode badge: always visible when bypass is on, even while
the pill row is collapsed (returns null when off). */}
<BypassPermissionsToggle />
{/* Permission-level pill: always visible, even while the pill row
is collapsed; opens the permission level dropdown. */}
<PermissionModeComposerPill side={effectiveMenuSide} />
{composerExpanded ? (
<>
<WebSearchToggle />
@ -2620,36 +2623,6 @@ const ArtifactsToggle: FC = () => {
);
};
// Claude gold pill shown while Bypass permissions is on; click to turn it off.
// Mirror of shared-composer's badge so both composers surface the state.
const BypassPermissionsToggle: FC = () => {
const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
const setBypassPermissions = useChatRuntimeStore(
(s) => s.setBypassPermissions,
);
if (!bypassPermissions) return null;
return (
<button
type="button"
onClick={() => setBypassPermissions(false)}
className="composer-pill-btn"
data-active="true"
data-variant="danger"
aria-label="Disable Bypass permissions"
title="Bypass permissions is on (no confirmation, no sandbox). Click to turn off."
>
<PillGlyph>
<HugeiconsIcon
icon={ShieldBanIcon}
strokeWidth={2}
className="size-[15px]"
/>
</PillGlyph>
<span>Bypass permissions</span>
</button>
);
};
const ToolStatusDisplay: FC = () => {
const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);

View file

@ -173,6 +173,7 @@ interface ResponseDetailsMetadata {
artifacts: boolean;
confirmToolCalls: boolean;
bypassPermissions: boolean;
permissionMode?: string;
};
}
@ -1454,6 +1455,11 @@ async function autoLoadSmallestModel(): Promise<{
blockedByTrustRemoteCode = true;
return false;
}
// Never install packages from a background load; explicit loads raise the upgrade dialog.
if (validation.requires_transformers_upgrade) {
hadNonTrustFailure = true;
return false;
}
return true;
}
@ -1946,6 +1952,7 @@ export function createOpenAIStreamAdapter(
mcpEnabledForChat,
confirmToolCalls,
bypassPermissions,
permissionMode,
webFetchToolsEnabled,
ragEnabled,
ragSource,
@ -2637,6 +2644,7 @@ export function createOpenAIStreamAdapter(
artifacts: renderHtmlToolEnabledForThisTurn,
confirmToolCalls,
bypassPermissions,
permissionMode,
},
});
const externalCapabilities = getProviderCapabilities(
@ -2948,6 +2956,16 @@ export function createOpenAIStreamAdapter(
...(supportsPreserveThinking
? { preserve_thinking: preserveThinking }
: {}),
// Permission level for local tool calls is sent for every local
// chat, not only when a tool pill is on: a process policy
// (unsloth run --enable-tools) can open the tool loop with no pill,
// and the backend must still see the selected gate. ask/auto request
// the confirm gate ("auto" only pauses calls flagged unsafe); off
// and full never prompt, full also drops the sandbox.
permission_mode: permissionMode,
confirm_tool_calls:
permissionMode === "ask" || permissionMode === "auto",
bypass_permissions: bypassPermissions,
...(supportsTools &&
(toolsEnabled ||
codeToolsEnabled ||
@ -2969,10 +2987,6 @@ export function createOpenAIStreamAdapter(
: []),
],
mcp_enabled: mcpEnabledForChat,
// Bypass Permissions wins: never request the confirm gate
// while bypassing, and tell the backend to drop the sandbox.
confirm_tool_calls: confirmToolCalls && !bypassPermissions,
bypass_permissions: bypassPermissions,
// Scope: thread_id = this thread's docs, kb_id = a KB,
// project_id = the thread's project sources (auto-on whenever
// the project has indexed sources, no Docs pill needed).

View file

@ -14,45 +14,49 @@ import {
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { DropdownMenuItem } from "@/components/ui/dropdown-menu";
import {
DropdownMenuSub,
DropdownMenuSubContent,
DropdownMenuSubTrigger,
} from "@/components/ui/dropdown-menu";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { Tick02Icon } from "@/lib/tick-icon";
import { PermissionModeMenuItems } from "./permission-mode-select";
// "Bypass permissions" entry for the composer "+" -> More menu. Mirrors the
// settings toggle: enabling demands the danger warning, disabling is immediate.
// The menu closes normally on select (no preventDefault) -- the warning dialog
// lives outside the menu (BypassPermissionsConfirmDialog, mounted once at the
// chat-page root and driven by the store), so it survives the menu unmounting
// and the "+"/More popovers don't stay frozen.
// "Bypass permissions" entry for the composer "+" -> More menu. Like the MCP
// pill, it opens a submenu where the user picks the permission level (Ask for
// approval / Approve for me / Full access). Picking Full access demands the
// danger warning; the other levels apply immediately. The menu closes normally
// on select (no preventDefault) -- the warning dialog lives outside the menu
// (BypassPermissionsConfirmDialog, mounted once at the chat-page root and
// driven by the store), so it survives the menu unmounting and the "+"/More
// popovers don't stay frozen.
export function BypassPermissionsMenuItem() {
const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
const setBypassPermissions = useChatRuntimeStore(
(s) => s.setBypassPermissions,
);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const setBypassConfirmOpen = useChatRuntimeStore(
(s) => s.setBypassConfirmOpen,
);
return (
<DropdownMenuItem
className={bypassPermissions ? "text-bypass font-medium" : undefined}
onSelect={() => {
if (bypassPermissions) {
setBypassPermissions(false);
} else {
// Defer past Radix's menu-close focus restoration: opening the dialog
// synchronously here lets the dropdown grab focus back and breaks the
// dialog's focus trap.
setTimeout(() => setBypassConfirmOpen(true), 0);
<DropdownMenuSub>
<DropdownMenuSubTrigger
className={
permissionMode === "full" ? "text-bypass font-medium" : undefined
}
}}
>
<HugeiconsIcon icon={ShieldBanIcon} strokeWidth={2} />
Bypass permissions
{bypassPermissions ? (
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
) : null}
</DropdownMenuItem>
>
<HugeiconsIcon icon={ShieldBanIcon} strokeWidth={2} />
Bypass permissions
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="unsloth-plus-menu w-[300px]">
<PermissionModeMenuItems
// Defer past Radix's menu-close focus restoration: opening the
// dialog synchronously here lets the dropdown grab focus back and
// breaks the dialog's focus trap.
onRequestFullAccess={() =>
setTimeout(() => setBypassConfirmOpen(true), 0)
}
/>
</DropdownMenuSubContent>
</DropdownMenuSub>
);
}
@ -63,19 +67,17 @@ export function BypassPermissionsMenuItem() {
export function BypassPermissionsConfirmDialog() {
const open = useChatRuntimeStore((s) => s.bypassConfirmOpen);
const setOpen = useChatRuntimeStore((s) => s.setBypassConfirmOpen);
const setBypassPermissions = useChatRuntimeStore(
(s) => s.setBypassPermissions,
);
const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode);
return (
<AlertDialog open={open} onOpenChange={setOpen}>
<AlertDialogContent size="sm">
<AlertDialogHeader>
<AlertDialogTitle>Enable Bypass permissions?</AlertDialogTitle>
<AlertDialogTitle>Enable Full access?</AlertDialogTitle>
<AlertDialogDescription>
Bypass permissions is dangerous since the AI model might delete,
corrupt your machine, and or cause real world damage to you or the
world - only accept if you are certain
Full access (Bypass permissions) is dangerous since the AI model
might delete, corrupt your machine, and or cause real world damage
to you or the world - only accept if you are certain
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
@ -84,7 +86,7 @@ export function BypassPermissionsConfirmDialog() {
variant="destructive"
className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
onClick={() => {
setBypassPermissions(true);
setPermissionMode("full");
setOpen(false);
}}
>

View file

@ -6,16 +6,6 @@ import {
AlertDescription,
AlertTitle,
} from "@/components/ui/alert";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { Checkbox } from "@/components/ui/checkbox";
import {
@ -81,6 +71,7 @@ import { Fragment, type ReactNode } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { toast } from "@/lib/toast";
import { OpenAICodeExecSection } from "./components/openai-code-exec-section";
import { PermissionModeDropdown } from "./permission-mode-select";
import { resyncInferenceStatusAfterServerModelChange } from "./hooks/use-chat-model-runtime";
import {
type ExternalProviderConfig,
@ -2037,9 +2028,8 @@ function NudgeToolCallsToggle() {
}
function ConfirmToolCallsToggle() {
const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls);
const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls);
const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
return (
<div className="flex items-center justify-between gap-3">
@ -2049,85 +2039,49 @@ function ConfirmToolCallsToggle() {
Confirm tool calls
</span>
<InfoHint>
When on, local Studio tool calls pause for your approval before they
run. Provider-hosted tools are not gated here.
When on, every local Unsloth tool call pauses for your approval
before it runs (the "Ask for approval" level). When off, tool calls
run without prompts inside the sandbox (the "Off" level).
Provider-hosted tools are not gated here.
</InfoHint>
</div>
{bypassPermissions ? (
{permissionMode === "full" ? (
<span className="text-[11px] text-muted-foreground">
Overridden by Bypass permissions
Overridden by Full access (Bypass permissions)
</span>
) : null}
</div>
<Switch
className="panel-switch"
checked={confirmToolCalls && !bypassPermissions}
checked={permissionMode === "ask"}
onCheckedChange={setConfirmToolCalls}
disabled={bypassPermissions}
disabled={permissionMode === "full"}
/>
</div>
);
}
function BypassPermissionsToggle() {
const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
const setBypassPermissions = useChatRuntimeStore(
(s) => s.setBypassPermissions,
);
const [dialogOpen, setDialogOpen] = useState(false);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
return (
<div className="flex flex-col gap-1.5">
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Bypass permissions
</span>
<InfoHint>
Dangerous. Runs every tool call with no confirmation and disables
the python/terminal sandbox. Environment secrets are stripped, but
code can still read files and credentials on your machine.
</InfoHint>
</div>
<Switch
className="panel-switch"
checked={bypassPermissions}
onCheckedChange={(next) => {
if (next) setDialogOpen(true);
else setBypassPermissions(false);
}}
/>
<div className="flex flex-col gap-2">
<div className="flex min-w-0 items-center gap-1.5">
<span className="whitespace-nowrap text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Bypass permissions
</span>
<InfoHint>
How Unsloth approves tool calls before they run. Full access is
dangerous: it disables confirmations and the code sandbox.
</InfoHint>
</div>
{bypassPermissions ? (
{/* Full width, styled like the panel selects/preset input. */}
<PermissionModeDropdown triggerClassName="h-9 w-full justify-between rounded-full border-0 bg-[var(--panel-input-surface)] px-3.5 text-[13px] font-medium text-nav-fg shadow-none hover:bg-[var(--panel-input-surface)]" />
{permissionMode === "full" ? (
<span className="text-[11px] text-bypass">
Tool calls run with no confirmation and no sandbox.
</span>
) : null}
<AlertDialog open={dialogOpen} onOpenChange={setDialogOpen}>
<AlertDialogContent size="sm">
<AlertDialogHeader>
<AlertDialogTitle>Enable Bypass permissions?</AlertDialogTitle>
<AlertDialogDescription>
Bypass permissions is dangerous since the AI model might delete,
corrupt your machine, and or cause real world damage to you or the
world - only accept if you are certain
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>Cancel</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
onClick={() => {
setBypassPermissions(true);
setDialogOpen(false);
}}
>
I understand
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}

View file

@ -4,6 +4,10 @@
import { createElement, useCallback, useRef, useState } from "react";
import { toast } from "@/lib/toast";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import {
confirmTransformersUpgradeIfNeeded,
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { consumeNativePathToken } from "@/features/native-intents/api";
import {
notifyNative,
@ -245,6 +249,10 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string {
return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`;
}
function getTransformersUpgradeRequiredMessage(modelName: string): string {
return `${modelName} was not loaded because it needs a newer transformers release that was not installed. Load it again to install it.`;
}
/**
* Reconcile the chat runtime store against `/api/inference/status`: refresh the
* models/loras catalogs and either re-pin the active checkpoint or clear the
@ -626,6 +634,30 @@ export function useChatModelRuntime() {
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
});
// Upgrade consent runs before the security dialogs; Accept installs and the load continues.
if (validation.requires_transformers_upgrade) {
const upgraded = await confirmTransformersUpgradeIfNeeded({
modelName: modelId,
upgrade: validation.transformers_upgrade,
// No installable release: custom-code models may fall back to the trust_remote_code gate below.
trustRemoteCodeFallback: validation.requires_trust_remote_code,
});
// The install unloads the previous model before the swap (even when
// the swap then fails), so any exit after this point must roll back.
// False for the custom-code fallback, which resolves without installing.
if (
useTransformersUpgradeDialogStore
.getState()
.consumeServerUnloadedChat()
&& currentCheckpoint
) {
previousWasUnloaded = true;
}
if (!upgraded) {
throw new Error(getTransformersUpgradeRequiredMessage(displayName));
}
}
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
// Open the consent dialog when the model needs custom-code consent or has a
// flagged unsafe file. Fires even when trustRemoteCode is preset on, since the
// worker requires a matching fingerprint that only the dialog produces.

View file

@ -17,6 +17,7 @@ export {
type Preset,
} from "./chat-settings-sheet";
export { useChatRuntimeStore } from "./stores/chat-runtime-store";
export { PermissionModeDropdown } from "./permission-mode-select";
export { useChatSearchStore } from "./stores/chat-search-store";
export { usePinnedChatsStore } from "./stores/pinned-chats-store";
export { useChatPreferencesStore } from "./stores/chat-preferences-store";

View file

@ -0,0 +1,338 @@
// 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 {
ChevronDown,
CircleAlert,
CircleOff,
Hand,
ShieldCheck,
XIcon,
} from "lucide-react";
import { useState } from "react";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuLabel,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
import { Tick02Icon } from "@/lib/tick-icon";
import { cn } from "@/lib/utils";
import { HugeiconsIcon } from "@hugeicons/react";
import {
type PermissionMode,
useChatRuntimeStore,
} from "./stores/chat-runtime-store";
/**
* Permission levels for the Bypass permissions dropdowns (General settings,
* chat settings sheet, composer "+" menu). Off sits last as the toggle that
* turns the feature off entirely.
*/
export const PERMISSION_MODE_OPTIONS: readonly {
value: PermissionMode;
label: string;
description: string;
icon: typeof Hand;
}[] = [
{
value: "ask",
label: "Ask for approval",
description: "Always ask before tool calls edit files or use the internet",
icon: Hand,
},
{
value: "auto",
label: "Approve for me",
description: "Only ask for actions detected as potentially unsafe",
icon: ShieldCheck,
},
{
value: "full",
label: "Full access",
description:
"Unrestricted: no approval prompts and the code sandbox is disabled",
icon: CircleAlert,
},
{
value: "off",
label: "Off",
description: "Turn off bypass permissions",
icon: CircleOff,
},
] as const;
export function permissionModeOption(mode: PermissionMode) {
return (
PERMISSION_MODE_OPTIONS.find((option) => option.value === mode) ??
PERMISSION_MODE_OPTIONS[0]
);
}
/** The option rows shared by every permission dropdown/submenu. Non-full
* levels apply directly; picking Full access must go through the caller's
* danger confirmation, so it's a separate callback. */
export function PermissionModeMenuItems({
onRequestFullAccess,
}: {
onRequestFullAccess: () => void;
}) {
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode);
return (
<>
{PERMISSION_MODE_OPTIONS.map((option) => (
<DropdownMenuItem
key={option.value}
onSelect={() => {
// Reselecting the active level toggles the feature off.
if (option.value === permissionMode) {
setPermissionMode("off");
} else if (option.value === "full") {
onRequestFullAccess();
} else {
setPermissionMode(option.value);
}
}}
className={cn(
"items-start gap-2 py-2",
permissionMode === option.value && "font-medium",
option.value === "full" &&
permissionMode === "full" &&
"text-bypass",
)}
>
<option.icon className="mt-0.5 size-4 shrink-0" strokeWidth={2} />
<span className="flex min-w-0 flex-1 flex-col gap-0.5">
<span className="text-[13px] leading-tight">{option.label}</span>
<span className="text-xs font-normal leading-snug text-muted-foreground">
{option.description}
</span>
</span>
{permissionMode === option.value ? (
<HugeiconsIcon
icon={Tick02Icon}
strokeWidth={2}
className="ml-auto mt-0.5 size-4 shrink-0"
/>
) : null}
</DropdownMenuItem>
))}
</>
);
}
/** Danger confirmation shown before Full access turns on. Self-contained so
* the dropdown works outside the chat page (e.g. the Settings dialog). */
export function FullAccessConfirmDialog({
open,
onOpenChange,
}: {
open: boolean;
onOpenChange: (open: boolean) => void;
}) {
const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode);
return (
<AlertDialog open={open} onOpenChange={onOpenChange}>
<AlertDialogContent size="sm">
<AlertDialogHeader>
<AlertDialogTitle>Enable Full access?</AlertDialogTitle>
<AlertDialogDescription>
Full access (Bypass permissions) is dangerous since the AI model
might delete, corrupt your machine, and or cause real world damage
to you or the world - only accept if you are certain
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>Cancel</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
onClick={() => {
setPermissionMode("full");
onOpenChange(false);
}}
>
I understand
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}
/**
* Select-style dropdown (like the MCP composer menu) for picking the
* permission level. Used in General settings and the chat settings sheet.
*/
export function PermissionModeDropdown({
side = "bottom",
align = "end",
triggerClassName,
}: {
side?: "top" | "bottom";
align?: "start" | "end";
triggerClassName?: string;
} = {}) {
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const [confirmOpen, setConfirmOpen] = useState(false);
const active = permissionModeOption(permissionMode);
const ActiveIcon = active.icon;
return (
<>
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button
variant="outline"
size="sm"
className={cn(
"gap-1.5",
triggerClassName,
// Last so a text color in triggerClassName cannot override it.
permissionMode === "full" &&
"text-bypass hover:text-bypass border-bypass/50",
)}
aria-label="Permission level for tool calls"
>
<ActiveIcon className="size-3.5 shrink-0" strokeWidth={2} />
<span className="min-w-0 flex-1 truncate text-left">
{active.label}
</span>
<ChevronDown className="size-3.5 shrink-0 opacity-60" />
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent
side={side}
align={align}
className="w-[300px]"
avoidCollisions={true}
>
<DropdownMenuLabel>
How should tool calls be approved?
</DropdownMenuLabel>
<PermissionModeMenuItems
// Defer past the menu-close focus restoration so the dialog's
// focus trap isn't broken by the dropdown grabbing focus back.
onRequestFullAccess={() =>
setTimeout(() => setConfirmOpen(true), 0)
}
/>
</DropdownMenuContent>
</DropdownMenu>
<FullAccessConfirmDialog
open={confirmOpen}
onOpenChange={setConfirmOpen}
/>
</>
);
}
/**
* Composer pill (mirrors the MCP pill) showing the current permission level
* in the chat box; clicking opens the level dropdown. Danger-styled while
* Full access is on. The Full access pick routes through the store-driven
* BypassPermissionsConfirmDialog mounted at the chat-page root, so the
* warning survives this menu unmounting.
*/
export function PermissionModeComposerPill({
side = "bottom",
}: {
side?: "top" | "bottom";
} = {}) {
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const setBypassConfirmOpen = useChatRuntimeStore(
(s) => s.setBypassConfirmOpen,
);
const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode);
const active = permissionModeOption(permissionMode);
const ActiveIcon = active.icon;
const fullAccess = permissionMode === "full";
// Off means the feature is off: no pill (re-enable via the "+" menu or
// settings, like the pre-levels bypass badge).
if (permissionMode === "off") return null;
return (
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<button
type="button"
className="composer-pill-btn composer-pill-permissions"
data-pill-label={active.label}
data-active={fullAccess ? "true" : "false"}
data-variant={fullAccess ? "danger" : undefined}
aria-label="Permission level for tool calls"
title={`${active.label}: ${active.description}`}
>
{/* The icon doubles as an off switch (mirrors the MCP pill): hover
swaps it to an X; clicking it turns bypass permissions Off (no
prompts, sandbox on) without opening the menu. In compact
icon-only mode the glyph is the whole button, so clicks fall
through and open the menu instead. */}
<span
role="button"
aria-label="Turn off bypass permissions"
tabIndex={-1}
onPointerDown={(e) => {
if (e.currentTarget.closest('[data-pill-compact="true"]')) {
return;
}
e.stopPropagation();
}}
onClick={(e) => {
if (e.currentTarget.closest('[data-pill-compact="true"]')) {
return;
}
e.stopPropagation();
setPermissionMode("off");
}}
className="composer-pill-glyph cursor-pointer"
>
<ActiveIcon className="size-[15px]" strokeWidth={2} />
<XIcon className="composer-pill-x" />
</span>
<span>{active.label}</span>
<HugeiconsIcon
icon={ChevronDownStandardIcon}
strokeWidth={1.5}
className="composer-pill-caret size-[15px]"
/>
</button>
</DropdownMenuTrigger>
<DropdownMenuContent
side={side}
align="start"
sideOffset={0}
avoidCollisions={true}
className="unsloth-plus-menu w-[300px]"
>
<DropdownMenuLabel>
How should tool calls be approved?
</DropdownMenuLabel>
<PermissionModeMenuItems
// Defer past the menu-close focus restoration (see PermissionModeDropdown).
onRequestFullAccess={() =>
setTimeout(() => setBypassConfirmOpen(true), 0)
}
/>
</DropdownMenuContent>
</DropdownMenu>
);
}

View file

@ -48,7 +48,6 @@ import {
Image03Icon,
McpServerIcon,
PencilRulerIcon,
ShieldBanIcon,
} from "@hugeicons/core-free-icons";
import { useNavigate } from "@tanstack/react-router";
import { HugeiconsIcon } from "@hugeicons/react";
@ -62,11 +61,16 @@ import {
import { listPromptEntries, type PromptEntry } from "./api/prompts-api";
import { McpComposerButton } from "./mcp-composer-button";
import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item";
import { PermissionModeComposerPill } from "./permission-mode-select";
import { reasoningCapsFromLoad } from "./lib/apply-inference-status-to-store";
import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button";
import { NewProjectDialog } from "./components/new-project-dialog";
import { useChatProjects } from "./hooks/use-chat-projects";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import {
confirmTransformersUpgradeIfNeeded,
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { loadModel, validateModel } from "./api/chat-api";
import {
parseExternalModelId,
@ -506,6 +510,7 @@ export function SharedComposer({
);
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const setMcpEnabledForChat = useChatRuntimeStore(
(s) => s.setMcpEnabledForChat,
@ -525,10 +530,6 @@ export function SharedComposer({
const setWebFetchToolsEnabled = useChatRuntimeStore(
(s) => s.setWebFetchToolsEnabled,
);
const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
const setBypassPermissions = useChatRuntimeStore(
(s) => s.setBypassPermissions,
);
const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled);
const setRagEnabled = useChatRuntimeStore((s) => s.setRagEnabled);
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
@ -681,9 +682,12 @@ export function SharedComposer({
const ragDisabled = modelLoaded && (isExternalModel || !supportsTools);
const showRagPill = !isExternalModel;
// Above 4 pills, collapse to icons only to cut clutter. Compare, Search and
// Code always show; the rest are conditional.
// Code always show; the permission pill shows in every mode except "off"
// (it renders null there); the rest are conditional.
const permissionPillVisible = permissionMode !== "off";
const pillsCompact =
3 +
(permissionPillVisible ? 1 : 0) +
(showImagePill ? 1 : 0) +
(showRagPill && ragEnabled && !ragDisabled ? 1 : 0) +
(showWebFetchPill ? 1 : 0) +
@ -929,6 +933,10 @@ export function SharedComposer({
return parts[parts.length - 1] || id;
}
// Set when an accepted transformers install unloaded the active model
// server-side; a later failure must then clear the stale checkpoint.
let upgradeUnloadedActive = false;
// Helper: load a model and update store checkpoint
async function ensureModelLoaded(
sel: CompareModelSelection,
@ -955,6 +963,31 @@ export function SharedComposer({
trust_remote_code: loadTrustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
});
// Upgrade dialog first (mirrors the primary load path).
if (validation.requires_transformers_upgrade) {
const upgraded = await confirmTransformersUpgradeIfNeeded({
modelName: sel.id,
upgrade: validation.transformers_upgrade,
// No installable release: custom-code models may fall back to the trust_remote_code gate below.
trustRemoteCodeFallback: validation.requires_trust_remote_code,
});
// The install unloads the active model before the swap (even when the
// swap then fails); if a later gate cancels or the load fails, the UI
// must stop pointing at that unloaded model.
if (
useTransformersUpgradeDialogStore
.getState()
.consumeServerUnloadedChat()
&& currentStore.params.checkpoint
) {
upgradeUnloadedActive = true;
}
if (!upgraded) {
throw new Error(
`${modelDisplayName(sel.id)} needs a newer transformers release to load.`,
);
}
}
if (
validation.requires_trust_remote_code ||
validation.requires_security_review
@ -990,6 +1023,7 @@ export function SharedComposer({
tensor_parallel: currentStore.tensorParallel,
});
saveSpeculativeType(specSettings.speculativeType);
upgradeUnloadedActive = false;
const store = useChatRuntimeStore.getState();
store.setCheckpoint(
resp.model,
@ -1097,6 +1131,11 @@ export function SharedComposer({
toast.success("Compare complete", { id: toastId, duration: 2000 });
} catch (err) {
compareStepSucceededRef.current = false;
// The install already unloaded the previously active model; drop the
// checkpoint so the UI does not keep pointing at an unloaded model.
if (upgradeUnloadedActive) {
useChatRuntimeStore.getState().clearCheckpoint();
}
toast.error("Compare failed", {
id: toastId,
description: err instanceof Error ? err.message : "Unknown error",
@ -1617,29 +1656,10 @@ export function SharedComposer({
</PillGlyph>
<span>Compare</span>
</button>
{/* Bypass sits immediately after Compare and ahead of every other
tool pill (Search, Code, ...) so the active danger state reads
first; only Compare outranks it. */}
{bypassPermissions && (
<button
type="button"
onClick={() => setBypassPermissions(false)}
className="composer-pill-btn"
data-active="true"
data-variant="danger"
aria-label="Disable Bypass permissions"
title="Bypass permissions is on (no confirmation, no sandbox). Click to turn off."
>
<PillGlyph>
<HugeiconsIcon
icon={ShieldBanIcon}
strokeWidth={2}
className="size-[15px]"
/>
</PillGlyph>
<span>Bypass permissions</span>
</button>
)}
{/* Permission-level pill sits immediately after Compare and ahead
of every other tool pill (Search, Code, ...) so the Full access
danger state reads first; only Compare outranks it. */}
<PermissionModeComposerPill side="top" />
<button
type="button"
disabled={searchDisabled}

View file

@ -45,6 +45,19 @@ export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY =
export const MODELS_FIT_ON_DEVICE_ONLY_KEY =
"unsloth_models_fit_on_device_only";
export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions";
export const CHAT_PERMISSION_MODE_KEY = "unsloth_chat_permission_mode";
/**
* Permission level for local tool calls:
* - "ask": always ask before every tool call runs.
* - "auto" ("Approve for me"): only ask for calls the backend detects as
* potentially unsafe; read-only calls run immediately. Sandbox stays on.
* - "off": never ask; tool calls run automatically inside the sandbox
* (the original default before permission levels existed).
* - "full" ("Full access"): no confirmations and the python/terminal sandbox
* is disabled. Session-only; never restored from storage.
*/
export type PermissionMode = "ask" | "auto" | "off" | "full";
export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY =
"unsloth_chat_web_fetch_tools_enabled";
export const CHAT_RAG_SOURCE_KEY = "unsloth_chat_rag_source";
@ -319,8 +332,10 @@ export function loadOptionalBool(key: string): boolean | null {
/**
* Resolve the web-search / code-execution pill state to apply when a model
* loads. Honors the user's persisted preference so a tool-capable model never
* re-enables a pill the user turned off; falls back to the model's capability
* only when no preference has been expressed.
* re-enables a pill the user turned off, and never re-disables one they turned
* on. When no preference has been expressed the pills stay off: tool execution
* is opt-in, so the person enables it with a click rather than a tool-capable
* model turning it on for them.
*/
export function resolveToolsEnabledOnLoad(supportsTools: boolean): {
toolsEnabled: boolean;
@ -328,8 +343,8 @@ export function resolveToolsEnabledOnLoad(supportsTools: boolean): {
} {
if (!supportsTools) return { toolsEnabled: false, codeToolsEnabled: false };
return {
toolsEnabled: loadOptionalBool(CHAT_TOOLS_ENABLED_KEY) ?? true,
codeToolsEnabled: loadOptionalBool(CHAT_CODE_TOOLS_ENABLED_KEY) ?? true,
toolsEnabled: loadOptionalBool(CHAT_TOOLS_ENABLED_KEY) ?? false,
codeToolsEnabled: loadOptionalBool(CHAT_CODE_TOOLS_ENABLED_KEY) ?? false,
};
}
@ -342,6 +357,37 @@ function saveBool(key: string, value: boolean): void {
}
}
/**
* "full" is intentionally not restorable: it disables the sandbox and every
* confirmation gate, so it must be re-enabled (through the warning dialog)
* each session. First run falls back to the legacy "Confirm tool calls"
* toggle so existing users keep their behavior (on -> ask, explicitly
* off -> "off", i.e. no prompts); fresh installs default to "auto".
*/
function loadPermissionMode(): PermissionMode {
if (!canUseStorage()) return "auto";
try {
const raw = localStorage.getItem(CHAT_PERMISSION_MODE_KEY);
if (raw === "ask" || raw === "auto" || raw === "off") return raw;
} catch {
// ignore
}
const legacyConfirm = loadOptionalBool(CHAT_CONFIRM_TOOL_CALLS_KEY);
if (legacyConfirm === null) return "auto";
return legacyConfirm ? "ask" : "off";
}
function savePermissionMode(mode: PermissionMode): void {
if (!canUseStorage() || mode === "full") return;
try {
localStorage.setItem(CHAT_PERMISSION_MODE_KEY, mode);
} catch {
// ignore
}
}
const INITIAL_PERMISSION_MODE: PermissionMode = loadPermissionMode();
function loadString(key: string, fallback: string): string {
if (!canUseStorage()) return fallback;
try {
@ -514,7 +560,11 @@ export function isPendingGguf(pending: PendingModelSelection | null): boolean {
* wrong file. */
export function pendingSelectionMatches(
pending: PendingModelSelection | null,
pick: { id: string; ggufVariant?: string | null; nativePathToken?: string | null },
pick: {
id: string;
ggufVariant?: string | null;
nativePathToken?: string | null;
},
): boolean {
return (
pending != null &&
@ -615,8 +665,15 @@ type ChatRuntimeStore = {
* Bypass Permissions: when on, tool calls run with no confirmation gate
* AND the python/terminal execution sandbox is disabled on the backend
* (secrets are still stripped). Takes precedence over confirmToolCalls.
* Kept in sync with permissionMode ("full" <=> true).
*/
bypassPermissions: boolean;
/**
* Permission level. Single source of truth for the bypass dropdowns;
* bypassPermissions and confirmToolCalls mirror it so legacy call sites
* keep working. "full" is session-only (never persisted).
*/
permissionMode: PermissionMode;
/** Whether the "Enable Bypass Permissions?" warning dialog is open. Lifted out
* of the composer menu so confirming/cancelling it doesn't leave the menu frozen. */
bypassConfirmOpen: boolean;
@ -759,6 +816,7 @@ type ChatRuntimeStore = {
setMcpEnabledForChat: (enabled: boolean) => void;
setConfirmToolCalls: (enabled: boolean) => void;
setBypassPermissions: (enabled: boolean) => void;
setPermissionMode: (mode: PermissionMode) => void;
setBypassConfirmOpen: (open: boolean) => void;
allowToolAlways: (sessionId: string, toolName: string) => void;
setToolConfirmation: (
@ -1081,11 +1139,15 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
false,
),
mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false),
confirmToolCalls: loadBool(CHAT_CONFIRM_TOOL_CALLS_KEY, false),
// Mirrors permissionMode (gate requested for ask/auto) so both controls
// agree on load.
confirmToolCalls:
INITIAL_PERMISSION_MODE === "ask" || INITIAL_PERMISSION_MODE === "auto",
// Never restore Bypass Permissions from storage: it disables the sandbox and
// the confirmation gate, so it must be re-enabled (through the warning
// dialog) each session rather than silently reactivating on reload.
bypassPermissions: false,
permissionMode: INITIAL_PERMISSION_MODE,
bypassConfirmOpen: false,
alwaysAllowToolsBySession: new Map<string, Set<string>>(),
toolConfirmations: {},
@ -1453,14 +1515,53 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
return { mcpEnabledForChat };
}),
setConfirmToolCalls: (confirmToolCalls) =>
set(() => {
set((state) => {
saveBool(CHAT_CONFIRM_TOOL_CALLS_KEY, confirmToolCalls);
return { confirmToolCalls };
// The legacy toggle is a view over the permission level: on -> "ask",
// off -> "off" (no prompts). While "full" is active the level is left
// alone (the toggle is disabled in the UI anyway).
if (state.permissionMode === "full") return { confirmToolCalls };
const permissionMode: PermissionMode = confirmToolCalls ? "ask" : "off";
savePermissionMode(permissionMode);
return { confirmToolCalls, permissionMode };
}),
setPermissionMode: (permissionMode) =>
set(() => {
// "full" is session-only (never persisted, see init); ask/auto/off
// persist and keep the legacy confirm toggle in sync (the gate is
// requested for both ask and auto).
savePermissionMode(permissionMode);
if (permissionMode === "full") {
// Full access sends confirm_tool_calls=false; keep the store flag in
// sync so response metadata does not report confirmations as enabled.
return { permissionMode, bypassPermissions: true, confirmToolCalls: false };
}
const confirmToolCalls =
permissionMode === "ask" || permissionMode === "auto";
saveBool(CHAT_CONFIRM_TOOL_CALLS_KEY, confirmToolCalls);
return { permissionMode, bypassPermissions: false, confirmToolCalls };
}),
setBypassPermissions: (bypassPermissions) =>
// Deliberately not persisted (see init): a reload must not silently keep
// the sandbox/confirmation bypass active without re-accepting the warning.
set(() => ({ bypassPermissions })),
// Turning bypass off returns to the last persisted ask/auto level.
set(() => {
if (bypassPermissions) {
// Full access never prompts; mirror confirm_tool_calls=false in the
// store so metadata does not report confirmations as enabled.
return {
bypassPermissions,
permissionMode: "full" as PermissionMode,
confirmToolCalls: false,
};
}
const permissionMode = loadPermissionMode();
return {
bypassPermissions,
permissionMode,
confirmToolCalls: permissionMode === "ask" || permissionMode === "auto",
};
}),
setBypassConfirmOpen: (bypassConfirmOpen) =>
set(() => ({ bypassConfirmOpen })),
allowToolAlways: (sessionId, toolName) =>

View file

@ -1,6 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { TransformersUpgradeInfo } from "@/features/transformers-upgrade";
export interface BackendModelDetails {
id: string;
name?: string | null;
@ -78,6 +80,10 @@ export interface ValidateModelResponse {
requires_security_review?: boolean;
/** Native context length from the local GGUF header; null until downloaded. */
context_length?: number | null;
/** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */
requires_transformers_upgrade?: boolean;
/** Set only when requires_transformers_upgrade. */
transformers_upgrade?: TransformersUpgradeInfo | null;
}
export interface GgufVariantDetail {
@ -343,6 +349,15 @@ export interface OpenAIChatCompletionsRequest {
mcp_enabled?: boolean;
/** Local models + enable_tools only. */
confirm_tool_calls?: boolean;
/**
* Local models + enable_tools only. Gate level for local tool calls: "ask"
* prompts on every call, "auto" prompts only on calls flagged unsafe, "off"
* never prompts, "full" never prompts and drops the sandbox. Unset behaves
* as "ask".
*/
permission_mode?: "ask" | "auto" | "off" | "full";
/** Local models + enable_tools only. Full-access escape hatch. */
bypass_permissions?: boolean;
/** `kb_id` is exclusive; otherwise project and thread scopes may combine. */
rag_scope?: {
kb_id?: string;

View file

@ -7,10 +7,10 @@ export function isBrowserOffline(): boolean {
const NETWORK_STATUS_EVENT = "unsloth-network-status";
const REMOTE_OFFLINE_TTL_MS = 30_000;
const HUGGING_FACE_REMOTE_ORIGINS = [
"https://huggingface.co",
"https://datasets-server.huggingface.co",
] as const;
// Discovery and repository pages are served by the main Hugging Face origin.
// Keep optional services such as datasets-server separate so an outage there
// cannot make the whole Hub appear offline.
const HUGGING_FACE_ORIGIN = "https://huggingface.co";
const noopUnsubscribe = () => undefined;
type RemoteNetworkScope = string | readonly string[];
@ -33,7 +33,7 @@ export function getBrowserOfflineRetryDelayMs(): number {
// recovery doesn't stall on platforms where navigator.onLine is stuck false.
return Math.max(
0,
getRemoteOfflineUntil(HUGGING_FACE_REMOTE_ORIGINS) - Date.now(),
getRemoteOfflineUntil(HUGGING_FACE_ORIGIN) - Date.now(),
);
}
@ -56,7 +56,7 @@ function getRemoteOfflineUntil(scope: RemoteNetworkScope): number {
}
export function isRemoteNetworkOffline(
scope: RemoteNetworkScope = HUGGING_FACE_REMOTE_ORIGINS,
scope: RemoteNetworkScope = HUGGING_FACE_ORIGIN,
): boolean {
return getRemoteOfflineUntil(scope) > Date.now();
}
@ -66,7 +66,7 @@ export function isHuggingFaceOffline(): boolean {
// WebKitGTK/Tauri webviews). The authoritative signal is the empirical
// remote-offline TTL, set when a real fetch fails and cleared on next success;
// navigator's online/offline events still drive re-evaluation.
return isRemoteNetworkOffline(HUGGING_FACE_REMOTE_ORIGINS);
return isRemoteNetworkOffline(HUGGING_FACE_ORIGIN);
}
export function markRemoteNetworkOnline(origin?: string): void {
@ -85,13 +85,13 @@ export function markRemoteNetworkOnline(origin?: string): void {
}
export function markRemoteNetworkOffline(
originOrTtl: string | number = HUGGING_FACE_REMOTE_ORIGINS[0],
originOrTtl: string | number = HUGGING_FACE_ORIGIN,
ttlMs = REMOTE_OFFLINE_TTL_MS,
): void {
const origin =
typeof originOrTtl === "string"
? originOrTtl
: HUGGING_FACE_REMOTE_ORIGINS[0];
: HUGGING_FACE_ORIGIN;
const ttl = typeof originOrTtl === "number" ? originOrTtl : ttlMs;
const nextUntil = Date.now() + ttl;
if (nextUntil <= (remoteOfflineUntilByOrigin.get(origin) ?? 0)) {

View file

@ -14,7 +14,7 @@ import { Input } from "@/components/ui/input";
import { Switch } from "@/components/ui/switch";
import { usePlatformStore } from "@/config/env";
import { resetOnboardingDone } from "@/features/auth";
import { useChatRuntimeStore } from "@/features/chat";
import { PermissionModeDropdown, useChatRuntimeStore } from "@/features/chat";
import { openModelsDir } from "@/features/native-intents";
import { emitTrainingRunsChanged } from "@/features/training";
import {
@ -80,6 +80,10 @@ const PREFS_KEYS: string[] = [
"unsloth_settings_active_tab",
// Chat runtime prefs
"unsloth_chat_auto_title",
"unsloth_chat_permission_mode",
// Legacy confirm key: loadPermissionMode falls back to it, so clear both or
// a reset would restore the old level instead of the fresh default.
"unsloth_chat_confirm_tool_calls",
"unsloth_hf_token",
"unsloth_auto_heal_tool_calls",
"unsloth_nudge_tool_calls",
@ -583,6 +587,15 @@ export function GeneralTab() {
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.general.permissions.sectionTitle")}>
<SettingsRow
label={t("settings.general.permissions.bypassLabel")}
description={t("settings.general.permissions.bypassDescription")}
>
<PermissionModeDropdown />
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.general.notifications.sectionTitle")}>
<SettingsRow
label={t("settings.general.notifications.showLlamaUpdates")}

View file

@ -0,0 +1,32 @@
// 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 { authFetch } from "@/features/auth";
import { readFastApiError } from "@/lib/format-fastapi-error";
interface InstallLatestTransformersResponse {
success: boolean;
version: string;
message: string;
/** The server unloaded the active chat model before the swap (set even on a
* structured failure, so callers can restore their model state). */
model_unloaded?: boolean;
/** On a version-mismatch failure: the release that superseded the requested
* one, so Retry can use it. */
latest_version?: string | null;
}
/** Consented install of the latest transformers into the sidecar; synchronous, can take minutes. */
export async function installLatestTransformers(
version: string,
): Promise<InstallLatestTransformersResponse> {
const response = await authFetch("/api/inference/install-latest-transformers", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ version }),
});
if (!response.ok) {
throw new Error(await readFastApiError(response));
}
return (await response.json()) as InstallLatestTransformersResponse;
}

View file

@ -0,0 +1,167 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Spinner } from "@/components/ui/spinner";
import { cn } from "@/lib/utils";
import { PackageIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store";
function modelDisplayName(modelName: string | null): string {
if (!modelName) return "This model";
return modelName.split("/").pop() || modelName;
}
/** Root-mounted consent dialog for models needing a newer transformers;
* Install runs the sidecar install and resumes the paused load on success. */
export function TransformersUpgradeDialog() {
const open = useTransformersUpgradeDialogStore((s) => s.open);
const modelName = useTransformersUpgradeDialogStore((s) => s.modelName);
const upgrade = useTransformersUpgradeDialogStore((s) => s.upgrade);
const phase = useTransformersUpgradeDialogStore((s) => s.phase);
const errorMessage = useTransformersUpgradeDialogStore((s) => s.errorMessage);
const trustRemoteCodeFallback = useTransformersUpgradeDialogStore(
(s) => s.trustRemoteCodeFallback,
);
const install = useTransformersUpgradeDialogStore((s) => s.install);
const resolve = useTransformersUpgradeDialogStore((s) => s.resolve);
const displayName = modelDisplayName(modelName);
const modelType = upgrade?.model_type ?? "unknown";
const version = upgrade?.pypi_version ?? null;
// Only released PyPI versions are installable; dev (main) builds are never offered.
const installable = Boolean(upgrade?.supported_in_pypi && version);
const devOnly = !installable && Boolean(upgrade?.supported_in_main);
const installing = phase === "installing";
return (
<AlertDialog
open={open}
onOpenChange={(next) => {
// Escape/overlay dismiss must not abandon an in-flight install.
if (!next && !installing) resolve(false);
}}
>
<AlertDialogContent className="max-w-lg">
<AlertDialogHeader className="min-w-0">
<div className="flex w-full min-w-0 items-start gap-3">
<div className="flex size-9 shrink-0 items-center justify-center rounded-full bg-blue-500/10 text-blue-600 dark:text-blue-400">
<HugeiconsIcon icon={PackageIcon} className="size-5" />
</div>
<div className="min-w-0 flex-1 space-y-3">
<div className="space-y-1">
<AlertDialogTitle>New model architecture</AlertDialogTitle>
<AlertDialogDescription>
<span className="font-medium text-foreground">
{displayName}
</span>{" "}
uses the{" "}
<span className="font-mono text-foreground">{modelType}</span>{" "}
architecture, which your installed transformers does not
support yet.{" "}
{installable ? (
<>
Install transformers{" "}
<span className="font-medium text-foreground">
{version}
</span>{" "}
from PyPI to load it. The install runs once and can take
a minute; loading continues automatically afterwards.
</>
) : devOnly ? (
<>
Even the latest transformers release on PyPI does not
support it yet: the architecture is only available on the
transformers development branch (main), and Studio does
not install development builds. Support arrives with the
next transformers release on PyPI.
</>
) : (
<>
No released transformers version supports it yet, so it
cannot be loaded.
</>
)}
{!installable && trustRemoteCodeFallback ? (
<>
{" "}
This model also ships its own modeling code; you can
continue and review enabling that custom code instead.
</>
) : null}
</AlertDialogDescription>
</div>
{phase === "error" && errorMessage ? (
<p className="rounded-lg border border-red-500/30 bg-red-500/10 px-3 py-2 text-xs text-red-600 dark:text-red-400">
{errorMessage}
</p>
) : null}
{installing ? (
<p className="flex items-center gap-2 text-xs text-muted-foreground">
<Spinner className="size-3.5" />
Installing transformers {version}... This can take a minute.
</p>
) : null}
</div>
</div>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={installing}>Cancel</AlertDialogCancel>
{installable ? (
<>
{phase === "error" && trustRemoteCodeFallback ? (
// Install failed but the model ships custom code: offer the
// caller's trust_remote_code gate instead of forcing a retry.
<AlertDialogAction
className="bg-transparent text-foreground hover:bg-accent"
onClick={() => resolve(true)}
>
Continue with custom code
</AlertDialogAction>
) : null}
<AlertDialogAction
disabled={installing}
className={cn(installing && "pointer-events-none")}
onClick={(event) => {
// Keep the dialog open; the store closes it on success.
event.preventDefault();
void install();
}}
>
{installing ? (
<>
<Spinner className="size-4" />
Installing...
</>
) : phase === "error" ? (
"Retry install"
) : (
`Install transformers ${version}`
)}
</AlertDialogAction>
</>
) : trustRemoteCodeFallback ? (
// No installable release but the model ships custom code: continue
// into the caller's trust_remote_code gate as the last resort.
<AlertDialogAction onClick={() => resolve(true)}>
Continue with custom code
</AlertDialogAction>
) : null}
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}

View file

@ -0,0 +1,28 @@
// 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 { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store";
import type { TransformersUpgradeInfo } from "../types";
interface ConfirmArgs {
modelName: string;
/** validate's transformers_upgrade payload; null/undefined skips the dialog. */
upgrade: TransformersUpgradeInfo | null | undefined;
/** When no release is installable, offer continuing into the caller's custom-code gate. */
trustRemoteCodeFallback?: boolean;
}
/** Pause a load needing a newer transformers on the consent dialog and run the install.
* Resolves true when the load can continue; false on cancel or not-installable with no fallback. */
export async function confirmTransformersUpgradeIfNeeded({
modelName,
upgrade,
trustRemoteCodeFallback,
}: ConfirmArgs): Promise<boolean> {
if (!upgrade) return true;
return useTransformersUpgradeDialogStore
.getState()
.requestConsent(modelName, upgrade, {
trustRemoteCodeFallback: Boolean(trustRemoteCodeFallback),
});
}

View file

@ -0,0 +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
export { TransformersUpgradeDialog } from "./components/transformers-upgrade-dialog";
export { confirmTransformersUpgradeIfNeeded } from "./hooks/use-transformers-upgrade-consent";
export { installLatestTransformers } from "./api/transformers-upgrade-api";
export { useTransformersUpgradeDialogStore } from "./stores/transformers-upgrade-dialog-store";
export type { TransformersUpgradeInfo } from "./types";

View file

@ -0,0 +1,139 @@
// 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 { create } from "zustand";
import { installLatestTransformers } from "../api/transformers-upgrade-api";
import type { TransformersUpgradeInfo, TransformersUpgradePhase } from "../types";
type Resolver = (installed: boolean) => void;
// One in-flight consent; a new request resolves any prior pending one as declined.
let pendingResolver: Resolver | null = null;
interface TransformersUpgradeDialogStore {
open: boolean;
modelName: string | null;
upgrade: TransformersUpgradeInfo | null;
phase: TransformersUpgradePhase;
errorMessage: string | null;
/** Model ships custom code; without a PyPI install the load may fall back to trust_remote_code. */
trustRemoteCodeFallback: boolean;
/** True once this consent's install completed. The install unloads the previous
* model before swapping, so the caller must treat it as already unloaded; the
* custom-code fallback resolves true without installing and leaves it loaded. */
installRan: boolean;
/** True when the server unloaded the active chat model during this consent,
* including a swap that failed AFTER the unload: callers must then treat
* their previous model as gone and roll back on any later cancel. */
serverUnloadedChat: boolean;
/** Read-and-clear serverUnloadedChat: each waiter consumes the signal once,
* so a superseding consent can neither erase it before the old waiter reads
* it nor leak it into an unrelated later load. */
consumeServerUnloadedChat: () => boolean;
/** Open the dialog for a paused load; resolves true on install success or custom-code fallback. */
requestConsent: (
modelName: string,
upgrade: TransformersUpgradeInfo,
options?: { trustRemoteCodeFallback?: boolean },
) => Promise<boolean>;
/** Accept/Retry: run the install; on success resolve(true) and close. */
install: () => Promise<void>;
resolve: (installed: boolean) => void;
}
export const useTransformersUpgradeDialogStore =
create<TransformersUpgradeDialogStore>()((set, get) => ({
open: false,
modelName: null,
upgrade: null,
phase: "consent",
errorMessage: null,
trustRemoteCodeFallback: false,
installRan: false,
serverUnloadedChat: false,
requestConsent: (modelName, upgrade, options) =>
new Promise<boolean>((resolve) => {
pendingResolver?.(false);
pendingResolver = resolve;
set({
open: true,
modelName,
upgrade,
phase: "consent",
errorMessage: null,
trustRemoteCodeFallback: Boolean(options?.trustRemoteCodeFallback),
installRan: false,
});
}),
consumeServerUnloadedChat: () => {
const value = get().serverUnloadedChat;
if (value) set({ serverUnloadedChat: false });
return value;
},
install: async () => {
const { upgrade, phase } = get();
const version = upgrade?.pypi_version;
if (!version || phase === "installing") return;
const requestResolver = pendingResolver;
set({ phase: "installing", errorMessage: null });
let result: Awaited<ReturnType<typeof installLatestTransformers>>;
try {
result = await installLatestTransformers(version);
// Latch the server-side unload IMMEDIATELY, before any resolver-identity
// guard: even a superseded consent's install may have unloaded the chat
// model, and the signal must survive for whichever load consumes it next.
if (result.model_unloaded) {
set({ serverUnloadedChat: true });
}
} catch (error) {
// Ignore the failure if a newer request superseded this consent.
if (pendingResolver === requestResolver) {
set({
phase: "error",
errorMessage:
error instanceof Error && error.message
? error.message
: "Failed to install transformers.",
});
}
return;
}
if (pendingResolver === requestResolver) {
if (result.success) {
// serverUnloadedChat was latched above (and is never reset here): a
// retry after a failed-after-unload attempt reports false because the
// model is already gone, and a superseded install may have set it too.
set({ installRan: true });
get().resolve(true);
return;
}
// Structured failure: the swap failed but may have already unloaded the
// chat model; record that so a later cancel still rolls the caller back.
// A version mismatch also carries the superseding release, so Retry
// re-requests a version that can actually succeed.
const { upgrade } = get();
set({
phase: "error",
errorMessage: result.message || "Failed to install transformers.",
serverUnloadedChat:
get().serverUnloadedChat || Boolean(result.model_unloaded),
...(result.latest_version && upgrade
? { upgrade: { ...upgrade, pypi_version: result.latest_version } }
: {}),
});
}
},
resolve: (installed) => {
const resolver = pendingResolver;
pendingResolver = null;
set({
open: false,
modelName: null,
upgrade: null,
phase: "consent",
errorMessage: null,
trustRemoteCodeFallback: false,
});
resolver?.(installed);
},
}));

View file

@ -0,0 +1,16 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/** Wire shape of `transformers_upgrade` from /api/inference/validate. */
export interface TransformersUpgradeInfo {
/** config.json model_type unknown to installed transformers. */
model_type: string;
/** Latest transformers release on PyPI at check time. */
pypi_version?: string | null;
/** Latest PyPI release ships this model_type (installable after consent). */
supported_in_pypi?: boolean;
/** Only transformers main ships it (dev-only; not installable). */
supported_in_main?: boolean;
}
export type TransformersUpgradePhase = "consent" | "installing" | "error";

View file

@ -181,6 +181,12 @@ export const en = {
revoked: "All preview links revoked",
revokeError: "Couldn't revoke preview links",
},
permissions: {
sectionTitle: "Permissions",
bypassLabel: "Bypass permissions",
bypassDescription:
"How Unsloth approves chat tool calls (terminal, python, web, MCP) before they run. Full access disables approvals and the code sandbox.",
},
notifications: {
sectionTitle: "Notifications",
showLlamaUpdates: "llama.cpp update notifications",

View file

@ -1483,6 +1483,15 @@ html[data-chat-font] .aui-root {
.composer-pill-btn[data-active="true"] {
color: var(--primary);
}
/* Permission-level pill: higher-contrast grey than the resting pills so
the active level stays legible (darker in light mode, lighter in dark).
Full access keeps the danger yellow below. */
.composer-pill-btn.composer-pill-permissions:not([data-variant="danger"]) {
color: color-mix(in oklab, var(--foreground) 60%, transparent);
}
.dark .composer-pill-btn.composer-pill-permissions:not([data-variant="danger"]) {
color: color-mix(in oklab, var(--foreground) 72%, transparent);
}
/* Bypass permissions badge: bright yellow text, no resting fill; the
rounded hover pill picks up the yellow accent like other toggles. */
.composer-pill-btn[data-variant="danger"] {

View file

@ -4086,7 +4086,9 @@ if ($script:StudioVtOk -and -not $env:NO_COLOR) {
}
Write-Host " $Rule" -ForegroundColor DarkGray
}
step "launch" "unsloth studio -H 0.0.0.0 -p 8888"
step "launch" "unsloth studio -p 8888"
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 ""
# Match studio/setup.sh: exit non-zero for degraded llama.cpp when called

View file

@ -1975,8 +1975,8 @@ else
else
printf " ${C_DIM}%-15s${C_OK}%s${C_RST}\n" "launch" "unsloth studio -p 8888"
fi
printf " ${C_DIM}%-15s%s${C_RST}\n" "" "(add -H 0.0.0.0 to allow network / cloud access)"
printf " ${C_DIM}%-15s%s${C_RST}\n" "" "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
printf " ${C_DIM}%-15s%s${C_RST}\n" "" "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
printf " ${C_DIM}%-15s%s${C_RST}\n" "" "(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)"
fi
echo ""

View file

@ -0,0 +1,201 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""CPU-only routing tests for the single-pass GGUF export and parallel quantization.
With convert/quantize monkeypatched, verify save_to_gguf's pass planning:
- a single directly-convertible output type (f32/f16/bf16/q8_0) converts in ONE pass
with no llama-quantize step and no 16-bit intermediate,
- k-quants and imatrix runs keep the two-pass route,
- multiple quantize passes run through the bounded pool with request order preserved,
- quantize failures still raise the actionable RuntimeError.
"""
from __future__ import annotations
import contextlib
import os
import threading
import time
import pytest
import unsloth.save as save_mod
# -- _choose_first_conversion (pure planning logic) ----------------------------------------
@pytest.mark.parametrize(
"methods, model_dtype, expected",
[
(["q8_0"], "f16", "q8_0"), # default "fast_quantized" path: single pass
(["q8_0", "q8_0"], "bf16", "q8_0"), # duplicates collapse to a single pass
(["f32"], "f16", "f32"), # 16/32-bit outputs convert directly too
(["bf16"], "bf16", "bf16"),
(["q4_k_m"], "f16", "f16"), # k-quants need a 16-bit base
(["q4_k_m", "q8_0"], "bf16", "bf16"), # mixes need the shared base
(["q8_0", "f16"], "f16", "f16"),
],
)
def test_choose_first_conversion(methods, model_dtype, expected):
assert save_mod._choose_first_conversion(methods, model_dtype) == expected
def test_choose_first_conversion_imatrix_forces_two_pass():
# Only llama-quantize can apply an imatrix, so q8_0-only must keep the 16-bit base.
assert save_mod._choose_first_conversion(["q8_0"], "f16", has_imatrix = True) == "f16"
# -- save_to_gguf pass planning (mocked convert/quantize) -----------------------------------
class _Harness:
"""Monkeypatched convert/quantize recording calls and creating real files."""
def __init__(
self,
monkeypatch,
tmp_path,
quantize_delays = None,
quantize_error = None,
):
self.tmp_path = tmp_path
self.convert_calls = []
self.quantize_calls = []
self.active = 0
self.max_concurrency = 0
self._lock = threading.Lock()
self._delays = quantize_delays or {}
self._error = quantize_error
monkeypatch.setattr(save_mod, "check_llama_cpp", lambda: ("llama-quantize", "convert.py"))
monkeypatch.setattr(
save_mod,
"_download_convert_hf_to_gguf",
lambda: (str(tmp_path / "convert.py"), {"LlamaForCausalLM"}, set()),
)
monkeypatch.setattr(save_mod, "use_local_gguf", contextlib.nullcontext)
monkeypatch.setattr(save_mod, "convert_to_gguf", self._convert)
monkeypatch.setattr(save_mod, "quantize_gguf", self._quantize)
def _convert(self, **kwargs):
self.convert_calls.append(kwargs)
suffix = kwargs["quantization_type"]
if suffix == "None":
suffix = kwargs["model_dtype"]
out = self.tmp_path / f"{kwargs['model_name']}.{suffix.upper()}.gguf"
out.write_bytes(b"GGUF")
return [str(out)], False
def _quantize(
self,
input_gguf,
output_gguf,
quant_type,
imatrix = None,
n_threads = None,
**kw,
):
with self._lock:
self.active += 1
self.max_concurrency = max(self.max_concurrency, self.active)
try:
if self._error is not None:
raise self._error
time.sleep(self._delays.get(quant_type, 0.02))
self.quantize_calls.append({"quant_type": quant_type, "n_threads": n_threads})
with open(output_gguf, "wb") as f:
f.write(b"GGUF")
return output_gguf
finally:
with self._lock:
self.active -= 1
def _run(tmp_path, methods, **kwargs):
model_dir = tmp_path / "model_dir"
model_dir.mkdir(exist_ok = True)
return save_mod.save_to_gguf(
model_name = "testmodel",
model_type = "llama",
model_dtype = "float16",
model_directory = str(model_dir),
quantization_method = methods,
**kwargs,
)
def test_q8_0_only_is_single_pass(monkeypatch, tmp_path):
h = _Harness(monkeypatch, tmp_path)
locations, want_full_precision, _ = _run(tmp_path, ["q8_0"])
assert len(h.convert_calls) == 1
assert h.convert_calls[0]["quantization_type"] == "q8_0"
assert h.quantize_calls == [], "single-pass export must not launch llama-quantize"
assert want_full_precision is True, "the converted file IS the requested output"
assert len(locations) == 1 and locations[0].endswith("testmodel.Q8_0.gguf")
assert os.path.exists(locations[0])
def test_fast_quantized_alias_is_single_pass(monkeypatch, tmp_path):
h = _Harness(monkeypatch, tmp_path)
_run(tmp_path, "fast_quantized") # the default of save_pretrained_gguf
assert h.convert_calls[0]["quantization_type"] == "q8_0"
assert h.quantize_calls == []
def test_k_quant_keeps_two_pass(monkeypatch, tmp_path):
h = _Harness(monkeypatch, tmp_path)
locations, want_full_precision, _ = _run(tmp_path, ["q4_k_m"])
assert h.convert_calls[0]["quantization_type"] == "f16"
assert [c["quant_type"] for c in h.quantize_calls] == ["q4_k_m"]
assert want_full_precision is False
# The 16-bit intermediate must be cleaned up.
assert len(locations) == 1 and locations[0].endswith("testmodel.Q4_K_M.gguf")
def test_mixed_methods_share_16bit_base(monkeypatch, tmp_path):
h = _Harness(monkeypatch, tmp_path)
_run(tmp_path, ["q4_k_m", "q8_0"])
assert h.convert_calls[0]["quantization_type"] == "f16"
assert sorted(c["quant_type"] for c in h.quantize_calls) == ["q4_k_m", "q8_0"]
def test_parallel_quants_preserve_request_order(monkeypatch, tmp_path):
# First method is the slowest: completion order != request order.
h = _Harness(
monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.3, "q5_k_m": 0.05, "q6_k": 0.01}
)
locations, _, _ = _run(tmp_path, ["q4_k_m", "q5_k_m", "q6_k"])
assert h.max_concurrency == 2, "quantize passes should overlap, bounded at 2"
quant_names = [os.path.basename(l) for l in locations if "F16" not in l]
assert quant_names == [
"testmodel.Q6_K.gguf", # list is reversed by the cleanup block, as before
"testmodel.Q5_K_M.gguf",
"testmodel.Q4_K_M.gguf",
]
assert all(
c["n_threads"] is not None for c in h.quantize_calls
), "parallel workers must split the thread budget explicitly"
def test_parallel_quants_env_kill_switch(monkeypatch, tmp_path):
monkeypatch.setenv("UNSLOTH_PARALLEL_GGUF_QUANTS", "0")
h = _Harness(monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.05, "q5_k_m": 0.05})
_run(tmp_path, ["q4_k_m", "q5_k_m"])
assert h.max_concurrency == 1
def test_duplicate_methods_quantize_once(monkeypatch, tmp_path):
h = _Harness(monkeypatch, tmp_path)
_run(tmp_path, ["q4_k_m", "q4_k_m"])
assert [c["quant_type"] for c in h.quantize_calls] == ["q4_k_m"]
def test_quantize_failure_raises_actionable_error(monkeypatch, tmp_path):
h = _Harness(monkeypatch, tmp_path, quantize_error = OSError("disk full"))
with pytest.raises(RuntimeError, match = "Quantization failed"):
_run(tmp_path, ["q4_k_m", "q5_k_m"])

View file

@ -102,6 +102,19 @@ class TestGetModelName(unittest.TestCase):
),
("unsloth/Kimi-K2-Instruct", True, "unsloth/Kimi-K2-Instruct-BF16", True),
("unsloth/Kimi-K2-Instruct", False, "unsloth/Kimi-K2-Instruct", False),
# DeepScaleR-1.5B must resolve to its own 16bit repo, not another model
(
"agentica-org/DeepScaleR-1.5B-Preview",
False,
"unsloth/DeepScaleR-1.5B-Preview",
True,
),
(
"agentica-org/DeepScaleR-1.5B-Preview",
True,
"unsloth/DeepScaleR-1.5B-Preview-unsloth-bnb-4bit",
True,
),
# Fallback-to-original behavior
"nonexistent-user/nonexistent-model-123",
"google/gemma-3-random-prototype-123",
@ -157,6 +170,10 @@ class TestGetModelName(unittest.TestCase):
with self.subTest(src = src):
self.assertEqual(FLOAT_TO_INT_MAPPER[src], expected)
self.assertEqual(MAP_TO_UNSLOTH_16bit["qwen/qwen3-8b-fp8"], "unsloth/Qwen3-8B-FP8")
self.assertEqual(
MAP_TO_UNSLOTH_16bit["agentica-org/deepscaler-1.5b-preview"],
"unsloth/DeepScaleR-1.5B-Preview",
)
if __name__ == "__main__":

View file

@ -1,4 +1,4 @@
"""Regression tests for the installers' post-install autostart opt-out."""
"""Regression tests for installer controls and process exits."""
from __future__ import annotations
@ -120,6 +120,27 @@ def test_windows_skip_autostart_bypasses_only_the_interactive_prompt():
)
@pytest.mark.skipif(shutil.which("pwsh") is None, reason = "PowerShell is unavailable")
def test_windows_installer_invalid_package_fails():
result = subprocess.run(
[
"pwsh",
"-NoProfile",
"-NonInteractive",
"-File",
str(INSTALL_PS1),
"--package",
"bad!",
],
capture_output = True,
text = True,
timeout = 30,
)
assert result.returncode != 0
assert "package name contains invalid characters" in result.stdout + result.stderr
def test_skip_autostart_is_documented_for_all_installers():
readme = README.read_text(encoding = "utf-8")
assert "UNSLOTH_SKIP_AUTOSTART=1 sh" in readme

View file

@ -141,6 +141,26 @@ def test_raw_text_loader():
except ValueError as e:
assert "stride" in str(e) and "chunk_size" in str(e)
# smart_chunk_text validation: called directly, chunk_size/stride are its own
# arguments and bypass the constructor guard, so it must guard itself or an
# invalid stride makes `start_idx += chunk_size - stride` non-positive and the
# chunking loop never terminates (hangs).
long_text = "This is a test file for raw text training. " * 10
valid_chunks = loader.smart_chunk_text(long_text, chunk_size = 5, stride = 2)
assert len(valid_chunks) > 0, "Valid stride should produce chunks"
try:
loader.smart_chunk_text(long_text, chunk_size = 5, stride = 5)
assert False, "Should raise ValueError for stride == chunk_size"
except ValueError as e:
assert "stride" in str(e) and "chunk_size" in str(e)
try:
loader.smart_chunk_text(long_text, chunk_size = 5, stride = 10)
assert False, "Should raise ValueError for stride > chunk_size"
except ValueError as e:
assert "stride" in str(e) and "chunk_size" in str(e)
# Preprocessor.
preprocessor = TextPreprocessor()
clean_text = preprocessor.clean_text(" messy text \n\n\n ")

View file

@ -132,6 +132,13 @@ class RawTextDataLoader:
3. Maintains context with stride overlap
4. Returns tokenized chunks directly (more efficient) or text chunks
"""
if chunk_size <= 0:
raise ValueError(f"chunk_size must be positive, got {chunk_size}")
if stride >= chunk_size:
raise ValueError(
f"stride ({stride}) must be smaller than chunk_size ({chunk_size}) to progress the chunking loop"
)
# Tokenize the whole text once for accurate token counts
tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False)
tokens = tokenized["input_ids"]

View file

@ -285,9 +285,6 @@ _cross_entropy_backward = triton.heuristics(
)(_cross_entropy_backward)
MAX_FUSED_SIZE = 65536 # 2**16
class Fast_CrossEntropyLoss(torch.autograd.Function):
@staticmethod
def forward(

View file

@ -270,7 +270,7 @@ try:
return fast_rms_layernorm(self, X, gemma = False)
except:
except (ImportError, AttributeError):
pass
@ -281,7 +281,7 @@ def patch_rms_layernorm():
try:
import transformers.models.mllama.modeling_mllama
transformers.models.mllama.modeling_mllama.MllamaTextRMSNorm = Unsloth_MllamaTextRMSNorm
except:
except (ImportError, AttributeError, NameError):
pass
return
@ -293,7 +293,7 @@ def unpatch_rms_layernorm():
try:
import transformers.models.mllama.modeling_mllama
transformers.models.mllama.modeling_mllama.MllamaTextRMSNorm = MllamaTextRMSNorm
except:
except (ImportError, AttributeError, NameError):
pass
return

View file

@ -701,7 +701,7 @@ __INT_TO_FLOAT_MAPPER = \
"unsloth/Qwen2.5-VL-72B-Instruct-bnb-4bit",
),
"unsloth/DeepScaleR-1.5B-Preview-unsloth-bnb-4bit" : (
"unsloth/DeepHermes-3-Llama-3-8B-Preview",
"unsloth/DeepScaleR-1.5B-Preview",
"agentica-org/DeepScaleR-1.5B-Preview",
"unsloth/DeepScaleR-1.5B-Preview-bnb-4bit",
),

View file

@ -739,7 +739,7 @@ class FastBaseModel:
if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1":
raise RuntimeError(
"Unsloth: UNSLOTH_VLLM_STANDBY is True, but UNSLOTH_VLLM_STANDBY is not set to 1!"
"Unsloth: `unsloth_vllm_standby` is True, but environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!"
)
if model_types is None:

View file

@ -1734,6 +1734,33 @@ def get_executable(executables):
return None
# Output types convert_hf_to_gguf.py can emit directly via --outtype.
_DIRECT_CONVERT_OUTTYPES = ("f32", "f16", "bf16", "q8_0")
def _choose_first_conversion(
quantization_methods,
model_dtype,
has_imatrix = False,
):
"""Pick the dtype of the initial HF -> GGUF conversion.
Single-pass fast path: when exactly one output type is requested and
convert_hf_to_gguf.py can emit it directly (f32/f16/bf16/q8_0), convert straight to
it - the llama-quantize pass and the 16-bit intermediate file are skipped entirely.
An imatrix forces the two-pass route since only llama-quantize can apply one.
Every other case converts to the source dtype first, so each requested method is
quantized from weights identical to the checkpoint's.
"""
unique_methods = set(quantization_methods)
if len(unique_methods) == 1 and not has_imatrix:
only_method = next(iter(unique_methods))
if only_method in _DIRECT_CONVERT_OUTTYPES:
return only_method
return model_dtype
def save_to_gguf(
model_name: str,
model_type: str,
@ -1782,10 +1809,6 @@ def save_to_gguf(
)
model_dtype = "f16"
# Check first_conversion as well
if first_conversion is None:
first_conversion = model_dtype
has_imatrix = imatrix is not None and str(imatrix) != ""
if has_imatrix:
# quantize_gguf gained the imatrix kwarg in a recent unsloth_zoo; fail fast (before the
@ -1834,32 +1857,12 @@ def save_to_gguf(
first_conversion = "None" # No quantization for GPT-OSS
# Only keep one conversion method since GPT-OSS doesn't quantize
quantization_method = ["None"]
else:
if first_conversion is None:
# Check if q8_0 is the ONLY quantization method requested
if len(quantization_method) == 1 and quantization_method[0] == "q8_0":
first_conversion = "None" # Let llama-quantize do the direct conversion
else:
# For all other cases, choose the highest precision format
# that can be requantized to all requested formats
strength = 0
for quant_method in quantization_method:
if quant_method == "f32":
strength = max(strength, 3)
elif quant_method == "f16":
strength = max(strength, 2)
elif quant_method == "bf16":
strength = max(strength, 1)
# Note: we don't set strength for q8_0 here since we handle it above
if strength >= 3:
first_conversion = "f32"
elif strength >= 2:
first_conversion = "f16"
elif strength >= 1:
first_conversion = "bf16"
else:
first_conversion = "bf16" # requantizing from q8_0 disallowed in new llama.cpp default to bf16.
elif first_conversion is None:
first_conversion = _choose_first_conversion(
quantization_method,
model_dtype,
has_imatrix = has_imatrix,
)
# Check bfloat16 support again for first_conversion
if first_conversion == "bf16" and not torch.cuda.is_bf16_supported():
@ -1868,12 +1871,19 @@ def save_to_gguf(
first_conversion_dtype = "" if first_conversion == "None" else first_conversion
# Print conversion info
needs_quantize_pass = any(m != first_conversion for m in quantization_method)
if needs_quantize_pass:
second_step = f"[2] Converting GGUF {first_conversion_dtype} to {quantization_method} might take 10 minutes each."
total_line = "In total, you will have to wait at least 16 minutes."
else:
second_step = f"[2] Single-pass export: converting straight to {quantization_method} - no separate quantize step."
total_line = "In total, you will have to wait at least 6 minutes."
print_info = (
f"==((====))== Unsloth: Conversion from HF to GGUF information\n"
f" {chr(92)}{chr(92)} /| [0] Installing llama.cpp might take 3 minutes.\n"
f"O^O/ {chr(92)}_/ {chr(92)} [1] Converting HF to GGUF {first_conversion_dtype} might take 3 minutes.\n"
f"{chr(92)} / [2] Converting GGUF {first_conversion_dtype} to {quantization_method} might take 10 minutes each.\n"
f' "-____-" In total, you will have to wait at least 16 minutes.\n'
f"{chr(92)} / {second_step}\n"
f' "-____-" {total_line}\n'
)
print(print_info)
@ -1959,75 +1969,154 @@ def save_to_gguf(
if not is_gpt_oss:
base_gguf = initial_files[0]
quants_created = False
for quant_method in quantization_method:
if quant_method != first_conversion:
# Deduplicate while keeping order; methods equal to the base conversion already
# exist on disk and need no quantize pass.
methods_to_quantize = [
m for m in dict.fromkeys(quantization_method) if m != first_conversion
]
def _quantize_one(quant_method, n_threads = None):
output_location = os.path.join(
gguf_directory, f"{model_name}.{quant_method.upper()}.gguf"
)
try:
if quant_method == "q2_k_l":
return _quantize_q2_k_l(
input_gguf = base_gguf,
output_gguf = output_location,
quantizer_location = quantizer_location,
n_threads = n_threads if n_threads is not None else n_cpus,
print_output = print_output,
imatrix = imatrix,
)
else:
# Use unsloth-zoo's standard quantization for all other methods. Only pass
# imatrix when set so older unsloth_zoo (no imatrix kwarg) still works for
# plain quants; an imatrix that cannot be applied was rejected above.
quant_kwargs = dict(
input_gguf = base_gguf,
output_gguf = output_location,
quant_type = quant_method,
quantizer_location = quantizer_location,
print_output = print_output,
)
if has_imatrix:
quant_kwargs["imatrix"] = imatrix
if n_threads is not None:
quant_kwargs["n_threads"] = n_threads
return quantize_gguf(**quant_kwargs)
except Exception as e:
if IS_KAGGLE_ENVIRONMENT:
raise RuntimeError(
f"Unsloth: Quantization failed for {output_location}\n"
"You are in a Kaggle environment, which might be the reason this is failing.\n"
"Kaggle only provides 20GB of disk space in the working directory.\n"
"Merging to 16bit for 7b models use 16GB of space.\n"
"This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"
"`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"
"You can try saving it to the `/tmp` directory for larger disk space.\n"
"I suggest you to save the 16bit model first, then use manual llama.cpp conversion.\n"
f"Error: {e}"
)
else:
if IS_WINDOWS:
build_instructions = (
f'cd "{LLAMA_CPP_DEFAULT_DIR}"\n'
f"cmake -S . -B build -DBUILD_SHARED_LIBS=OFF\n"
f"cmake --build build --config Release"
)
else:
build_instructions = (
f'cd "{LLAMA_CPP_DEFAULT_DIR}" && make clean && make all -j'
)
raise RuntimeError(
f"Unsloth: Quantization failed for {output_location}\n"
"You might have to compile llama.cpp yourself, then run this again.\n"
"You do not need to close this Python program. Run the following commands in a new terminal:\n"
f'git clone --recursive https://github.com/ggerganov/llama.cpp "{LLAMA_CPP_DEFAULT_DIR}"\n'
f"{build_instructions}\n"
"Once that's done, redo the quantization.\n"
f"Error: {e}"
)
# Outputs already on disk pre-date this run; never delete them on a failure.
preexisting_outputs = {
m
for m in methods_to_quantize
if os.path.exists(os.path.join(gguf_directory, f"{model_name}.{m.upper()}.gguf"))
}
# Each llama-quantize pass loads the whole base GGUF into RAM, so only run two at
# once when the host has headroom for two copies, else a multi-quant export that
# fit sequentially could OOM.
try:
base_bytes = sum(
os.path.getsize(f)
for f in initial_files
if "-mmproj" not in os.path.basename(f).lower()
)
mem_ok = psutil.virtual_memory().available >= int(2.5 * base_bytes)
except Exception:
mem_ok = False
# Independent llama-quantize runs on the same base GGUF can overlap. Kept at 2
# workers; run sequentially when streaming logs (UNSLOTH_ENABLE_LOGGING), on
# Kaggle/Colab, when RAM is tight, or when the kill switch (0/false/no/off/empty)
# is set.
_parallel_flag = os.environ.get("UNSLOTH_PARALLEL_GGUF_QUANTS", "1").strip().lower()
parallel_quants = (
len(methods_to_quantize) > 1
and not print_output
and not IS_KAGGLE_ENVIRONMENT
and not IS_COLAB_ENVIRONMENT
and mem_ok
and _parallel_flag not in ("0", "false", "no", "off", "")
)
if parallel_quants:
max_workers = min(2, len(methods_to_quantize))
# Split the thread budget so total threads match the sequential run.
per_worker_threads = max(1, n_cpus // max_workers)
print(
f"Unsloth: [2] Converting GGUF {first_conversion_dtype} into "
f"{methods_to_quantize}, {max_workers} at a time. This might take 10 minutes each..."
)
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_EXCEPTION
quantized_files = [None] * len(methods_to_quantize)
with ThreadPoolExecutor(max_workers = max_workers) as pool:
future_to_idx = {
pool.submit(_quantize_one, method, per_worker_threads): i
for i, method in enumerate(methods_to_quantize)
}
done, pending = wait(future_to_idx, return_when = FIRST_EXCEPTION)
# Do not start queued passes after a failure (avoid filling the disk).
for fut in pending:
fut.cancel()
first_exc = next((f.exception() for f in done if f.exception() is not None), None)
if first_exc is not None:
# Remove only outputs this run newly created; a file that pre-dated the
# run (or a canceled pass that never wrote) is left intact, so a rerun
# never deletes a prior artifact. Base kept for retry.
wait(future_to_idx)
for method in methods_to_quantize:
if method in preexisting_outputs:
continue
Path(
os.path.join(gguf_directory, f"{model_name}.{method.upper()}.gguf")
).unlink(missing_ok = True)
raise first_exc
for fut, i in future_to_idx.items():
quantized_files[i] = fut.result()
else:
quantized_files = []
for quant_method in methods_to_quantize:
print(
f"Unsloth: [2] Converting GGUF {first_conversion_dtype} into {quant_method}. This might take 10 minutes..."
)
output_location = os.path.join(
gguf_directory, f"{model_name}.{quant_method.upper()}.gguf"
)
try:
if quant_method == "q2_k_l":
quantized_file = _quantize_q2_k_l(
input_gguf = base_gguf,
output_gguf = output_location,
quantizer_location = quantizer_location,
n_threads = n_cpus,
print_output = print_output,
imatrix = imatrix,
)
else:
# Use unsloth-zoo's standard quantization for all other methods. Only pass
# imatrix when set so older unsloth_zoo (no imatrix kwarg) still works for
# plain quants; an imatrix that cannot be applied was rejected above.
quant_kwargs = dict(
input_gguf = base_gguf,
output_gguf = output_location,
quant_type = quant_method,
quantizer_location = quantizer_location,
print_output = print_output,
)
if has_imatrix:
quant_kwargs["imatrix"] = imatrix
quantized_file = quantize_gguf(**quant_kwargs)
all_saved_locations.append(quantized_file)
quants_created = True
except Exception as e:
if IS_KAGGLE_ENVIRONMENT:
raise RuntimeError(
f"Unsloth: Quantization failed for {output_location}\n"
"You are in a Kaggle environment, which might be the reason this is failing.\n"
"Kaggle only provides 20GB of disk space in the working directory.\n"
"Merging to 16bit for 7b models use 16GB of space.\n"
"This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"
"`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"
"You can try saving it to the `/tmp` directory for larger disk space.\n"
"I suggest you to save the 16bit model first, then use manual llama.cpp conversion.\n"
f"Error: {e}"
)
else:
if IS_WINDOWS:
build_instructions = (
f'cd "{LLAMA_CPP_DEFAULT_DIR}"\n'
f"cmake -S . -B build -DBUILD_SHARED_LIBS=OFF\n"
f"cmake --build build --config Release"
)
else:
build_instructions = (
f'cd "{LLAMA_CPP_DEFAULT_DIR}" && make clean && make all -j'
)
quantized_files.append(_quantize_one(quant_method))
raise RuntimeError(
f"Unsloth: Quantization failed for {output_location}\n"
"You might have to compile llama.cpp yourself, then run this again.\n"
"You do not need to close this Python program. Run the following commands in a new terminal:\n"
f'git clone --recursive https://github.com/ggerganov/llama.cpp "{LLAMA_CPP_DEFAULT_DIR}"\n'
f"{build_instructions}\n"
"Once that's done, redo the quantization.\n"
f"Error: {e}"
)
all_saved_locations.extend(quantized_files)
quants_created = len(quantized_files) > 0
print("Unsloth: Model files cleanup...")
want_full_precision = first_conversion in quantization_method
if quants_created:

View file

@ -0,0 +1,275 @@
You are a coding agent running in the Codex CLI, a terminal-based coding assistant. Codex CLI is an open source project led by OpenAI. You are expected to be precise, safe, and helpful.
Your capabilities:
- Receive user prompts and other context provided by the harness, such as files in the workspace.
- Communicate with the user by streaming thinking & responses, and by making & updating plans.
- Emit function calls to run terminal commands and apply patches. Depending on how this specific run is configured, you can request that these function calls be escalated to the user for approval before running. More on this in the "Sandbox and approvals" section.
Within this context, Codex refers to the open-source agentic coding interface (not the old Codex language model built by OpenAI).
# How you work
## Personality
Your default personality and tone is concise, direct, and friendly. You communicate efficiently, always keeping the user clearly informed about ongoing actions without unnecessary detail. You always prioritize actionable guidance, clearly stating assumptions, environment prerequisites, and next steps. Unless explicitly asked, you avoid excessively verbose explanations about your work.
# AGENTS.md spec
- Repos often contain AGENTS.md files. These files can appear anywhere within the repository.
- These files are a way for humans to give you (the agent) instructions or tips for working within the container.
- Some examples might be: coding conventions, info about how code is organized, or instructions for how to run or test code.
- Instructions in AGENTS.md files:
- The scope of an AGENTS.md file is the entire directory tree rooted at the folder that contains it.
- For every file you touch in the final patch, you must obey instructions in any AGENTS.md file whose scope includes that file.
- Instructions about code style, structure, naming, etc. apply only to code within the AGENTS.md file's scope, unless the file states otherwise.
- More-deeply-nested AGENTS.md files take precedence in the case of conflicting instructions.
- Direct system/developer/user instructions (as part of a prompt) take precedence over AGENTS.md instructions.
- The contents of the AGENTS.md file at the root of the repo and any directories from the CWD up to the root are included with the developer message and don't need to be re-read. When working in a subdirectory of CWD, or a directory outside the CWD, check for any AGENTS.md files that may be applicable.
## Responsiveness
### Preamble messages
Before making tool calls, send a brief preamble to the user explaining what youre about to do. When sending preamble messages, follow these principles and examples:
- **Logically group related actions**: if youre about to run several related commands, describe them together in one preamble rather than sending a separate note for each.
- **Keep it concise**: be no more than 1-2 sentences, focused on immediate, tangible next steps. (812 words for quick updates).
- **Build on prior context**: if this is not your first tool call, use the preamble message to connect the dots with whats been done so far and create a sense of momentum and clarity for the user to understand your next actions.
- **Keep your tone light, friendly and curious**: add small touches of personality in preambles feel collaborative and engaging.
- **Exception**: Avoid adding a preamble for every trivial read (e.g., `cat` a single file) unless its part of a larger grouped action.
**Examples:**
- “Ive explored the repo; now checking the API route definitions.”
- “Next, Ill patch the config and update the related tests.”
- “Im about to scaffold the CLI commands and helper functions.”
- “Ok cool, so Ive wrapped my head around the repo. Now digging into the API routes.”
- “Configs looking tidy. Next up is patching helpers to keep things in sync.”
- “Finished poking at the DB gateway. I will now chase down error handling.”
- “Alright, build pipeline order is interesting. Checking how it reports failures.”
- “Spotted a clever caching util; now hunting where it gets used.”
## Planning
You have access to an `update_plan` tool which tracks steps and progress and renders them to the user. Using the tool helps demonstrate that you've understood the task and convey how you're approaching it. Plans can help to make complex, ambiguous, or multi-phase work clearer and more collaborative for the user. A good plan should break the task into meaningful, logically ordered steps that are easy to verify as you go.
Note that plans are not for padding out simple work with filler steps or stating the obvious. The content of your plan should not involve doing anything that you aren't capable of doing (i.e. don't try to test things that you can't test). Do not use plans for simple or single-step queries that you can just do or answer immediately.
Do not repeat the full contents of the plan after an `update_plan` call — the harness already displays it. Instead, summarize the change made and highlight any important context or next step.
Before running a command, consider whether or not you have completed the previous step, and make sure to mark it as completed before moving on to the next step. It may be the case that you complete all steps in your plan after a single pass of implementation. If this is the case, you can simply mark all the planned steps as completed. Sometimes, you may need to change plans in the middle of a task: call `update_plan` with the updated plan and make sure to provide an `explanation` of the rationale when doing so.
Use a plan when:
- The task is non-trivial and will require multiple actions over a long time horizon.
- There are logical phases or dependencies where sequencing matters.
- The work has ambiguity that benefits from outlining high-level goals.
- You want intermediate checkpoints for feedback and validation.
- When the user asked you to do more than one thing in a single prompt
- The user has asked you to use the plan tool (aka "TODOs")
- You generate additional steps while working, and plan to do them before yielding to the user
### Examples
**High-quality plans**
Example 1:
1. Add CLI entry with file args
2. Parse Markdown via CommonMark library
3. Apply semantic HTML template
4. Handle code blocks, images, links
5. Add error handling for invalid files
Example 2:
1. Define CSS variables for colors
2. Add toggle with localStorage state
3. Refactor components to use variables
4. Verify all views for readability
5. Add smooth theme-change transition
Example 3:
1. Set up Node.js + WebSocket server
2. Add join/leave broadcast events
3. Implement messaging with timestamps
4. Add usernames + mention highlighting
5. Persist messages in lightweight DB
6. Add typing indicators + unread count
**Low-quality plans**
Example 1:
1. Create CLI tool
2. Add Markdown parser
3. Convert to HTML
Example 2:
1. Add dark mode toggle
2. Save preference
3. Make styles look good
Example 3:
1. Create single-file HTML game
2. Run quick sanity check
3. Summarize usage instructions
If you need to write a plan, only write high quality plans, not low quality ones.
## Task execution
You are a coding agent. Please keep going until the query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. Autonomously resolve the query to the best of your ability, using the tools available to you, before coming back to the user. Do NOT guess or make up an answer.
You MUST adhere to the following criteria when solving queries:
- Working on the repo(s) in the current environment is allowed, even if they are proprietary.
- Analyzing code for vulnerabilities is allowed.
- Showing user code and tool call details is allowed.
- Use the `apply_patch` tool to edit files (NEVER try `applypatch` or `apply-patch`, only `apply_patch`): {"command":["apply_patch","*** Begin Patch\\n*** Update File: path/to/file.py\\n@@ def example():\\n- pass\\n+ return 123\\n*** End Patch"]}
If completing the user's task requires writing or modifying files, your code and final answer should follow these coding guidelines, though user instructions (i.e. AGENTS.md) may override these guidelines:
- Fix the problem at the root cause rather than applying surface-level patches, when possible.
- Avoid unneeded complexity in your solution.
- Do not attempt to fix unrelated bugs or broken tests. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
- Update documentation as necessary.
- Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task.
- Use `git log` and `git blame` to search the history of the codebase if additional context is required.
- NEVER add copyright or license headers unless specifically requested.
- Do not waste tokens by re-reading files after calling `apply_patch` on them. The tool call will fail if it didn't work. The same goes for making folders, deleting folders, etc.
- Do not `git commit` your changes or create new git branches unless explicitly requested.
- Do not add inline comments within code unless explicitly requested.
- Do not use one-letter variable names unless explicitly requested.
- NEVER output inline citations like "【F:README.md†L5-L14】" in your outputs. The CLI is not able to render these so they will just be broken in the UI. Instead, if you output valid filepaths, users will be able to click on them to open the files in their editor.
## Validating your work
If the codebase has tests or the ability to build or run, consider using them to verify that your work is complete.
When testing, your philosophy should be to start as specific as possible to the code you changed so that you can catch issues efficiently, then make your way to broader tests as you build confidence. If there's no test for the code you changed, and if the adjacent patterns in the codebases show that there's a logical place for you to add a test, you may do so. However, do not add tests to codebases with no tests.
Similarly, once you're confident in correctness, you can suggest or use formatting commands to ensure that your code is well formatted. If there are issues you can iterate up to 3 times to get formatting right, but if you still can't manage it's better to save the user time and present them a correct solution where you call out the formatting in your final message. If the codebase does not have a formatter configured, do not add one.
For all of testing, running, building, and formatting, do not attempt to fix unrelated bugs. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
Be mindful of whether to run validation commands proactively. In the absence of behavioral guidance:
- When running in the non-interactive approval mode **never**, proactively run tests, lint and do whatever you need to ensure you've completed the task.
- When working in interactive approval modes like **untrusted**, or **on-request**, hold off on running tests or lint commands until the user is ready for you to finalize your output, because these commands take time to run and slow down iteration. Instead suggest what you want to do next, and let the user confirm first.
- When working on test-related tasks, such as adding tests, fixing tests, or reproducing a bug to verify behavior, you may proactively run tests regardless of approval mode. Use your judgement to decide whether this is a test-related task.
## Ambition vs. precision
For tasks that have no prior context (i.e. the user is starting something brand new), you should feel free to be ambitious and demonstrate creativity with your implementation.
If you're operating in an existing codebase, you should make sure you do exactly what the user asks with surgical precision. Treat the surrounding codebase with respect, and don't overstep (i.e. changing filenames or variables unnecessarily). You should balance being sufficiently ambitious and proactive when completing tasks of this nature.
You should use judicious initiative to decide on the right level of detail and complexity to deliver based on the user's needs. This means showing good judgment that you're capable of doing the right extras without gold-plating. This might be demonstrated by high-value, creative touches when scope of the task is vague; while being surgical and targeted when scope is tightly specified.
## Sharing progress updates
For especially longer tasks that you work on (i.e. requiring many tool calls, or a plan with multiple steps), you should provide progress updates back to the user at reasonable intervals. These updates should be structured as a concise sentence or two (no more than 8-10 words long) recapping progress so far in plain language: this update demonstrates your understanding of what needs to be done, progress so far (i.e. files explores, subtasks complete), and where you're going next.
Before doing large chunks of work that may incur latency as experienced by the user (i.e. writing a new file), you should send a concise message to the user with an update indicating what you're about to do to ensure they know what you're spending time on. Don't start editing or writing large files before informing the user what you are doing and why.
The messages you send before tool calls should describe what is immediately about to be done next in very concise language. If there was previous work done, this preamble message should also include a note about the work done so far to bring the user along.
## Presenting your work and final message
Your final message should read naturally, like an update from a concise teammate. For casual conversation, brainstorming tasks, or quick questions from the user, respond in a friendly, conversational tone. You should ask questions, suggest ideas, and adapt to the users style. If you've finished a large amount of work, when describing what you've done to the user, you should follow the final answer formatting guidelines to communicate substantive changes. You don't need to add structured formatting for one-word answers, greetings, or purely conversational exchanges.
You can skip heavy formatting for single, simple actions or confirmations. In these cases, respond in plain sentences with any relevant next step or quick option. Reserve multi-section structured responses for results that need grouping or explanation.
The user is working on the same computer as you, and has access to your work. As such there's no need to show the full contents of large files you have already written unless the user explicitly asks for them. Similarly, if you've created or modified files using `apply_patch`, there's no need to tell users to "save the file" or "copy the code into a file"—just reference the file path.
If there's something that you think you could help with as a logical next step, concisely ask the user if they want you to do so. Good examples of this are running tests, committing changes, or building out the next logical component. If theres something that you couldn't do (even with approval) but that the user might want to do (such as verifying changes by running the app), include those instructions succinctly.
Brevity is very important as a default. You should be very concise (i.e. no more than 10 lines), but can relax this requirement for tasks where additional detail and comprehensiveness is important for the user's understanding.
### Final answer structure and style guidelines
You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value.
**Section Headers**
- Use only when they improve clarity — they are not mandatory for every answer.
- Choose descriptive names that fit the content
- Keep headers short (13 words) and in `**Title Case**`. Always start headers with `**` and end with `**`
- Leave no blank line before the first bullet under a header.
- Section headers should only be used where they genuinely improve scanability; avoid fragmenting the answer.
**Bullets**
- Use `-` followed by a space for every bullet.
- Merge related points when possible; avoid a bullet for every trivial detail.
- Keep bullets to one line unless breaking for clarity is unavoidable.
- Group into short lists (46 bullets) ordered by importance.
- Use consistent keyword phrasing and formatting across sections.
**Monospace**
- Wrap all commands, file paths, env vars, and code identifiers in backticks (`` `...` ``).
- Apply to inline examples and to bullet keywords if the keyword itself is a literal file/command.
- Never mix monospace and bold markers; choose one based on whether its a keyword (`**`) or inline code/path (`` ` ``).
**File References**
When referencing files in your response, make sure to include the relevant start line and always follow the below rules:
* Use inline code to make file paths clickable.
* Each reference should have a stand alone path. Even if it's the same file.
* Accepted: absolute, workspacerelative, a/ or b/ diff prefixes, or bare filename/suffix.
* Line/column (1based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1).
* Do not use URIs like file://, vscode://, or https://.
* Do not provide range of lines
* Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5
**Structure**
- Place related bullets together; dont mix unrelated concepts in the same section.
- Order sections from general → specific → supporting info.
- For subsections (e.g., “Binaries” under “Rust Workspace”), introduce with a bolded keyword bullet, then list items under it.
- Match structure to complexity:
- Multi-part or detailed results → use clear headers and grouped bullets.
- Simple results → minimal headers, possibly just a short list or paragraph.
**Tone**
- Keep the voice collaborative and natural, like a coding partner handing off work.
- Be concise and factual — no filler or conversational commentary and avoid unnecessary repetition
- Use present tense and active voice (e.g., “Runs tests” not “This will run tests”).
- Keep descriptions self-contained; dont refer to “above” or “below”.
- Use parallel structure in lists for consistency.
**Dont**
- Dont use literal words “bold” or “monospace” in the content.
- Dont nest bullets or create deep hierarchies.
- Dont output ANSI escape codes directly — the CLI renderer applies them.
- Dont cram unrelated keywords into a single bullet; split for clarity.
- Dont let keyword lists run long — wrap or reformat for scanability.
Generally, ensure your final answers adapt their shape and depth to the request. For example, answers to code explanations should have a precise, structured explanation with code references that answer the question directly. For tasks with a simple implementation, lead with the outcome and supplement only with whats needed for clarity. Larger changes can be presented as a logical walkthrough of your approach, grouping related steps, explaining rationale where it adds value, and highlighting next actions to accelerate the user. Your answers should provide the right level of detail while being easily scannable.
For casual greetings, acknowledgements, or other one-off conversational messages that are not delivering substantive information or structured results, respond naturally without section headers or bullet formatting.
# Tool Guidelines
## Shell commands
When using the shell, you must adhere to the following guidelines:
- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.)
- Do not use python scripts to attempt to output larger chunks of a file.
## `update_plan`
A tool named `update_plan` is available to you. You can use it to keep an uptodate, stepbystep plan for the task.
To create a new plan, call `update_plan` with a short list of 1sentence steps (no more than 5-7 words each) with a `status` for each step (`pending`, `in_progress`, or `completed`).
When steps have been completed, use `update_plan` to mark each finished step as `completed` and the next step you are working on as `in_progress`. There should always be exactly one `in_progress` step until everything is done. You can mark multiple items as complete in a single `update_plan` call.
If all steps are complete, ensure you call `update_plan` to mark all steps as `completed`.

View file

@ -0,0 +1,238 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Masked terminal password prompt for the first-exposure password change.
Mirror of ``studio/backend/auth/terminal_prompt.py`` -- keep the two in sync.
The CLI parent cannot import the backend package outside the studio venv, so the
reader is duplicated here (like the auth mirroring in ``commands/studio.py``).
Input echoes one ``*`` per character (unlike ``getpass``). All output goes to
stderr so redirected stdout stays clean.
"""
from __future__ import annotations
import os
import sys
from typing import Callable, TextIO
# Keep in sync with studio/backend/models/auth.py ChangePasswordRequest
# (new_password min_length) and studio/backend/auth/storage.py.
MIN_PASSWORD_LENGTH = 8
# Env var that supplies the initial admin password non-interactively (mirror in
# studio/backend/auth/terminal_prompt.py). Keep the name in sync.
SUPPLIED_PASSWORD_ENV = "UNSLOTH_STUDIO_PASSWORD"
_BACKSPACE_CHARS = ("\x7f", "\x08")
_SUBMIT_CHARS = ("\r", "\n")
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
def _read_masked_posix(prompt: str, out: TextIO) -> str:
import codecs
import termios
import tty
fd = sys.stdin.fileno()
old_attrs = termios.tcgetattr(fd)
out.write(prompt)
out.flush()
chars: list[str] = []
try:
with _RestoreTtyOnSignals(fd, old_attrs):
# cbreak + ISIG off (mirrors terminal_prompt.py): with ISIG on,
# Ctrl-Z would suspend mid-read and leave the shell no-echo before
# the finally restores it. Ctrl-C/Ctrl-Z arrive as \x03/\x1a here.
tty.setcbreak(fd)
new_attrs = termios.tcgetattr(fd)
new_attrs[3] &= ~termios.ISIG
termios.tcsetattr(fd, termios.TCSADRAIN, new_attrs)
# Decode byte-at-a-time with errors="replace" (mirrors
# terminal_prompt.py): text-mode read(1) can raise UnicodeDecodeError
# on a pasted non-UTF-8 password or yield a lone surrogate that later
# crashes pbkdf2. os.read + incremental decoder maps bad bytes to
# U+FFFD and continues.
decoder = codecs.getincrementaldecoder(sys.stdin.encoding or "utf-8")("replace")
submitted = False
while not submitted:
raw = os.read(fd, 1)
if not raw: # stream ended mid-line: abort, don't submit
raise EOFError
# One byte can complete >1 char, so iterate over the decoder's output.
for ch in decoder.decode(raw):
if ch in _SUBMIT_CHARS:
submitted = True
break
if ch == "\x03": # Ctrl-C (ISIG off: surfaces as a char)
raise KeyboardInterrupt
if ch in ("\x04", "\x1a"): # Ctrl-D / Ctrl-Z
if not chars:
raise EOFError
continue
if ch in _BACKSPACE_CHARS:
if chars:
chars.pop()
out.write("\b \b")
out.flush()
continue
if ch < " ": # other control characters
continue
chars.append(ch)
out.write("*")
out.flush()
finally:
termios.tcsetattr(fd, termios.TCSADRAIN, old_attrs)
out.write("\n")
out.flush()
return "".join(chars)
def _read_masked_windows(prompt: str, out: TextIO) -> str:
import msvcrt
out.write(prompt)
out.flush()
chars: list[str] = []
try:
while True:
ch = msvcrt.getwch()
if ch in _SUBMIT_CHARS:
break
if ch == "\x03": # Ctrl-C: getwch swallows the signal, re-raise
raise KeyboardInterrupt
if ch in ("\x04", "\x1a"): # Ctrl-D / Ctrl-Z
if not chars:
raise EOFError
continue
if ch in ("\x00", "\xe0"): # function/arrow key: swallow the code
msvcrt.getwch()
continue
if ch in _BACKSPACE_CHARS:
if chars:
chars.pop()
out.write("\b \b")
out.flush()
continue
if ch < " ":
continue
chars.append(ch)
out.write("*")
out.flush()
finally:
out.write("\n")
out.flush()
return "".join(chars)
def read_masked(prompt: str, out: TextIO | None = None) -> str:
"""Read one line with ``*`` echo. Raises KeyboardInterrupt on Ctrl-C and
EOFError on Ctrl-D/Ctrl-Z at an empty prompt."""
if out is None:
out = sys.stderr
if os.name == "nt":
return _read_masked_windows(prompt, out)
return _read_masked_posix(prompt, out)
def prompt_new_password(verify_current: Callable[[str], bool], out: TextIO | None = None) -> str:
"""Prompt for a new admin password until a valid, confirmed one is given.
``verify_current`` returns True when the candidate equals the current stored
password; such candidates are rejected. KeyboardInterrupt/EOFError propagate
so the caller can abort the launch.
"""
if out is None:
out = sys.stderr
while True:
password = read_masked("New password: ", out)
if len(password) < MIN_PASSWORD_LENGTH:
out.write(f"Password must be at least {MIN_PASSWORD_LENGTH} characters. Try again.\n")
out.flush()
continue
if verify_current(password):
out.write("New password must differ from the current password. Try again.\n")
out.flush()
continue
confirmation = read_masked("Confirm new password: ", out)
if confirmation != password:
out.write("Passwords do not match. Try again.\n")
out.flush()
continue
return password
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 backend 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
def validate_new_password(candidate: str, verify_current: Callable[[str], bool]) -> "str | None":
"""Error message if ``candidate`` is unacceptable (too short or equal to the
current password), else None. Same policy as the interactive loop."""
if len(candidate) < MIN_PASSWORD_LENGTH:
return f"Password must be at least {MIN_PASSWORD_LENGTH} characters."
if verify_current(candidate):
return "New password must differ from the current password."
return None

View file

@ -170,6 +170,43 @@ def _hermes_install_hint() -> str:
return _HERMES_WINDOWS_INSTALL_HINT if os.name == "nt" else _HERMES_POSIX_INSTALL_HINT
def _hermes_resume_oneshot_args(args: list[str]) -> list[str]:
"""Route resumed one-shot prompts through Hermes' session-aware chat command."""
has_resume = any(
arg in ("--resume", "-r", "--continue", "-c")
or arg.startswith(("--resume=", "--continue="))
or (len(arg) > 2 and arg.startswith(("-r", "-c")))
for arg in args
)
if not has_resume:
return args
rewritten = list(args)
for index, arg in enumerate(rewritten):
if arg in ("-z", "--oneshot"):
rewritten[index] = "-q"
elif len(arg) > 2 and arg.startswith("-z"):
# argparse accepts attached short-option values (`-zPROMPT` and
# `-z=PROMPT`); preserve the value byte-for-byte when switching to -q.
rewritten[index] = f"-q{arg[2:]}"
elif arg.startswith("--oneshot="):
rewritten[index] = f"--query={arg.partition('=')[2]}"
else:
continue
if any(item == "--usage-file" or item.startswith("--usage-file=") for item in args):
raise typer.BadParameter(
"Hermes cannot resume a one-shot session with --usage-file; remove that option."
)
prefix = ["chat", "-Q"]
if "--yolo" not in rewritten:
prefix.append("--yolo")
if "--accept-hooks" not in rewritten:
prefix.append("--accept-hooks")
rewritten = prefix + rewritten
return rewritten
return args
class LoadOptions(NamedTuple):
"""Model-load knobs forwarded to /api/inference/load when --model triggers a load."""
@ -840,6 +877,60 @@ def _merge_codex_config(existing: str, base: str) -> str:
)
# Keep custom-model behavior aligned with Codex's own unknown-model fallback. This
# Apache-2.0 prompt is copied from openai/codex rust-v0.144.0 models-manager/prompt.md.
_CODEX_FALLBACK_PROMPT = Path(__file__).parent.parent / "codex_fallback_prompt.md"
_CODEX_MODEL_CATALOG_MIN_VERSION = (0, 110, 0)
def _codex_supports_model_catalog() -> bool:
executable = shutil.which("codex")
if executable is None:
# A --no-launch recipe may be copied to another machine; assume a current Codex.
return True
try:
output = subprocess.check_output(
[executable, "--version"], text = True, timeout = 10, stderr = subprocess.DEVNULL
)
except Exception:
return False
match = re.search(r"(\d+)\.(\d+)\.(\d+)", output)
return bool(match) and tuple(int(part) for part in match.groups()) >= (
_CODEX_MODEL_CATALOG_MIN_VERSION
)
def _codex_model_catalog(model: dict) -> dict:
"""Return conservative metadata for a Studio model unknown to Codex's built-in catalog."""
model_id = model["id"]
window = model.get("context_length") or model.get("max_context_length")
entry = {
"slug": model_id,
"display_name": model_id,
"description": "Model served by Unsloth Studio",
"supported_reasoning_levels": [],
"shell_type": "default",
"visibility": "none",
"supported_in_api": True,
"priority": 99,
"availability_nux": None,
"upgrade": None,
"base_instructions": _CODEX_FALLBACK_PROMPT.read_text(encoding = "utf-8"),
"supports_reasoning_summaries": False,
"supports_reasoning_summary_parameter": False,
"support_verbosity": False,
"default_verbosity": None,
"apply_patch_tool_type": None,
"truncation_policy": {"mode": "bytes", "limit": 10_000},
"supports_parallel_tool_calls": False,
"experimental_supported_tools": [],
}
if window:
entry["context_window"] = int(window)
entry["max_context_window"] = int(window)
return {"models": [entry]}
def write_codex_config(base: str, model: dict, home: Path) -> None:
home.mkdir(parents = True, exist_ok = True)
@ -857,6 +948,16 @@ def write_codex_config(base: str, model: dict, home: Path) -> None:
f'model_provider = "{_CODEX_PROFILE}"\n'
f"model = {json.dumps(model['id'])}\n"
)
if _codex_supports_model_catalog() and _CODEX_FALLBACK_PROMPT.is_file():
catalog = home / "model-catalog.json"
catalog_text = json.dumps(_codex_model_catalog(model), indent = 2) + "\n"
if not catalog.exists() or catalog.read_text(encoding = "utf-8") != catalog_text:
catalog.write_text(catalog_text, encoding = "utf-8")
typer.echo(f"Updated {catalog}")
# Resolve relative to the profile file. This also survives WSL launching a Windows
# Codex binary, where a Linux absolute path inside TOML would not be usable.
profile_text += f"model_catalog_json = {json.dumps(catalog.name)}\n"
window = model.get("context_length") or model.get("max_context_length")
if window:
profile_text += f"model_context_window = {int(window)}\n"
@ -875,6 +976,16 @@ def _wsl_windows_executable(command: list) -> Optional[str]:
return None
def _wsl_windows_path(path: Path) -> str:
try:
translated = subprocess.check_output(["wslpath", "-w", str(path)], text = True).strip()
except (OSError, subprocess.CalledProcessError) as exc:
_fail(f"Could not translate WSL path {path}: {exc}")
if not translated:
_fail(f"Could not translate WSL path {path}")
return translated
def _looks_like_path(value: str) -> bool:
# A var only wants the WSLENV /p flag if its value is a filesystem path: an
# absolute POSIX path (/...), a UNC path (\\...), or a drive-qualified Windows
@ -1184,6 +1295,7 @@ def write_openclaw_config(
model: dict,
path: Path,
yolo: bool = False,
workspace_path: Optional[str] = None,
) -> None:
config = _read_json_object(path)
if config is None:
@ -1208,8 +1320,23 @@ def write_openclaw_config(
"models": [provider_model],
}
# Pin a default model, else OpenClaw drops into its setup agent ("no models available").
defaults = _subdict(_subdict(config, "agents"), "defaults")
agents = _subdict(config, "agents")
defaults = _subdict(agents, "defaults")
_subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}"
# OPENCLAW_STATE_DIR does not relocate the workspace. Keep it beside the managed
# config so ephemeral launches avoid ~/.openclaw and persisted sessions retain it.
workspace = path.parent / "workspace"
workspace.mkdir(parents = True, exist_ok = True, mode = 0o700)
defaults["workspace"] = workspace_path or str(workspace)
# Per-agent paths override agents.defaults.workspace and OPENCLAW_STATE_DIR. This
# config is itself an isolated Unsloth copy, so remove stale explicit paths and let
# OpenClaw resolve every listed agent beneath the managed defaults/state directory.
agent_list = agents.get("list")
if isinstance(agent_list, list):
for agent_config in agent_list:
if isinstance(agent_config, dict):
agent_config.pop("workspace", None)
agent_config.pop("agentDir", None)
# Unauthenticated loopback gateway: without auth.mode=none the client won't open
# the websocket. The daemon must still be started separately (`openclaw gateway`).
gateway = _subdict(config, "gateway")
@ -1339,9 +1466,11 @@ def write_opencode_config(
tools = ("edit", "bash", "webfetch")
if yolo:
# OpenCode has no --yolo flag; auto-approve is the config `permission` block
# (singular). Allow the prompting tools so tool calls don't block on the TUI. This
# rides inline (OPENCODE_CONFIG_CONTENT) so --yolo works even over a project config.
# (singular). Allow the prompting tools and paths outside the launch directory so
# tool calls don't block on the TUI. This rides inline (OPENCODE_CONFIG_CONTENT) so
# --yolo works even over a project config.
session_permission = {t: "allow" for t in tools}
session_permission["external_directory"] = {"*": "allow"}
config["permission"] = dict(session_permission)
else:
# Undo only what --yolo wrote: our yolo sets an explicit per-tool "allow" for these
@ -1358,6 +1487,8 @@ def write_opencode_config(
for tool in tools:
if permission.get(tool) == "allow":
permission[tool] = "ask"
if permission.get("external_directory") == {"*": "allow"}:
permission["external_directory"] = {"*": "ask"}
if json.dumps(config, sort_keys = True) != before:
_write_private_json(path, config)
typer.echo(f"Updated {path}")
@ -1629,8 +1760,18 @@ def openclaw(
)
with _session_config("openclaw", launch, persist = persist) as cfg:
config_path = cfg / "openclaw.json"
workspace_path = None
if _wsl_windows_executable(command):
workspace_path = _wsl_windows_path(cfg / "workspace")
# key lives in the config, not the env; --yolo writes the exec policy here too.
write_openclaw_config(base, key, entry, config_path, yolo = yolo)
write_openclaw_config(
base,
key,
entry,
config_path,
yolo = yolo,
workspace_path = workspace_path,
)
# Scope both config and state so OpenClaw never touches the user's ~/.openclaw.
env = {"OPENCLAW_CONFIG_PATH": str(config_path), "OPENCLAW_STATE_DIR": str(cfg)}
_run(base, entry, env, command, launch = launch, install_hint = install_hint)
@ -1729,6 +1870,8 @@ def hermes(
persist: bool = _PERSIST_OPTION,
):
"""Point Hermes (Nous Research) at the running Studio server and start it."""
native_args = [*_yolo_command_flags("hermes", yolo), *ctx.args]
command = ["hermes", *_hermes_resume_oneshot_args(native_args)]
base, key, entry = _connect(
api_key,
model,
@ -1736,7 +1879,6 @@ def hermes(
serve = serve,
launch = launch,
)
command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args]
install_hint = _hermes_install_hint()
with _session_config("hermes", launch, persist = persist) as home:
# HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state)

View file

@ -3,6 +3,7 @@
import importlib.util
import hashlib
import hmac
import json
import os
import platform
@ -21,6 +22,8 @@ from pathlib import Path
from typing import List, Optional
import typer
from unsloth_cli.commands import _password_prompt
studio_app = typer.Typer(help = "Unsloth Studio commands.")
@ -480,6 +483,14 @@ def _connect_auth_db() -> sqlite3.Connection:
auth_dir = STUDIO_HOME / "auth"
auth_dir.mkdir(parents = True, exist_ok = True)
conn = sqlite3.connect(auth_dir / "auth.db")
# Mirror backend storage.get_connection: this path can create auth/ and
# auth.db (the pre-exposure gate writes here first), and sqlite3.connect
# makes the DB 0644 under a 022 umask. Keep both private.
for _path, _mode in ((auth_dir, 0o700), (auth_dir / "auth.db", 0o600)):
try:
os.chmod(_path, _mode)
except OSError:
pass
conn.execute(
"""
CREATE TABLE IF NOT EXISTS auth_user (
@ -627,6 +638,526 @@ def _create_desktop_secret_in_cli() -> str:
conn.close()
def _should_prompt_password_change(
*, cloudflare: Optional[bool], host: str, secure: bool, api_only: bool
) -> bool:
"""Whether this launch will expose Studio through the Cloudflare tunnel.
CLI mirror of run.py's _cloudflare_tunnel_should_start, minus the Colab
case (Colab launches never come through this CLI path). --secure implies
the tunnel; --cloudflare only tunnels non-api-only wildcard binds.
"""
if secure:
return True
if cloudflare is not True:
return False
return host in ("0.0.0.0", "::") and not api_only
def _prompt_streams_interactive() -> bool:
"""The prompt needs a real terminal for input and for the masked echo."""
try:
return sys.stdin.isatty() and sys.stderr.isatty()
except (AttributeError, ValueError):
return False
def _bootstrap_deadline_active() -> bool:
"""Whether the backend's bootstrap shutdown deadline will arm.
Mirror of studio/backend/auth/bootstrap_timeout.py bootstrap_timeout_seconds:
unset/blank/malformed UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT falls back to the 1h
default (a typo must not remove protection); 0 or negative disables it.
"""
raw = os.environ.get("UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT", "").strip()
if not raw:
return True
try:
return int(raw) > 0
except ValueError:
return True
def _cli_update_password(conn: sqlite3.Connection, username: str, new_password: str) -> None:
"""CLI mirror of backend update_password + change-password route effects.
One transaction: rehash, rotate the JWT secret, clear must_change_password,
revoke refresh tokens (PR #6651 finding), and drop the desktop secret. File
cleanup happens after commit; a failed unlink must not roll the change back.
"""
password_salt, password_hash = _hash_password(new_password)
with conn:
conn.execute(
"""
UPDATE auth_user
SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
WHERE username = ?
""",
(password_salt, password_hash, secrets.token_urlsafe(64), username),
)
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.execute(
"DELETE FROM app_secrets WHERE key IN (?, ?)",
(DESKTOP_SECRET_HASH_KEY, DESKTOP_SECRET_CREATED_AT_KEY),
)
for stale in (BOOTSTRAP_PASSWORD_FILE, DESKTOP_SECRET_FILE):
stale_path = STUDIO_HOME / "auth" / stale
try:
stale_path.unlink(missing_ok = True)
except OSError as exc:
# The hash is already committed, so a failed unlink must NOT roll the
# change back. But a locked-yet-writable file (Windows AV, read-only
# auth dir) must be truncated: otherwise its stale plaintext survives
# and generate_bootstrap_password() would re-validate this revoked
# credential after a later reset-password deletes auth.db. Mirrors
# backend clear_bootstrap_password().
try:
stale_path.write_text("")
cleared = True
except OSError:
cleared = False
if cleared:
typer.echo(
f"Warning: could not remove stale {stale} file ({exc}); cleared its "
"contents so the old credential cannot be reused.",
err = True,
)
else:
typer.echo(
f"Warning: could not remove or clear stale {stale} file ({exc}); the "
"old credential is still on disk. Remove it manually to prevent reuse "
"after a reset.",
err = True,
)
def _apply_supplied_password_before_launch(supplied_password: "str | None") -> None:
"""Non-interactively set the INITIAL admin password (from --password /
UNSLOTH_STUDIO_PASSWORD / stdin) before the server binds, while the account
still has its auto-generated bootstrap password.
Only ever sets the FIRST password: an already-set one is a hard error (an
override would be an auth bypass on a public launch), and an invalid value
fails closed. Runs in the parent before any re-exec so the secret never
crosses to the child argv.
"""
if not supplied_password:
return
try:
conn = _connect_auth_db()
except (OSError, sqlite3.Error) as exc:
typer.echo(
f"Error: --password could not open the Studio auth database ({exc}); not starting.",
err = True,
)
raise typer.Exit(1)
try:
_ensure_cli_default_admin(conn)
conn.commit()
row = conn.execute(
"SELECT password_salt, password_hash, must_change_password "
"FROM auth_user WHERE username = ?",
(DEFAULT_ADMIN_USERNAME,),
).fetchone()
if not row:
typer.echo(
"Error: --password could not initialize the admin account; not starting.",
err = True,
)
raise typer.Exit(1)
if not row[2]:
typer.echo(
"Error: a Studio admin password is already set; --password only sets "
"the initial password. Run `unsloth studio reset-password` first "
"(or change it in the UI).",
err = True,
)
raise typer.Exit(1)
password_salt, password_hash = row[0], row[1]
def _is_current_password(candidate: str) -> bool:
return hmac.compare_digest(
_pbkdf2_hex(candidate, password_salt.encode("utf-8")), password_hash
)
problem = _password_prompt.validate_new_password(supplied_password, _is_current_password)
if problem is not None:
typer.echo(f"Error: {problem} Not starting.", err = True)
raise typer.Exit(1)
_cli_update_password(conn, DEFAULT_ADMIN_USERNAME, supplied_password)
typer.echo(f"Password updated for '{DEFAULT_ADMIN_USERNAME}'.", err = True)
except (OSError, sqlite3.Error) as exc:
# Any DB failure fails closed (typer.Exit is not caught here, so the
# deliberate Exit(1) branches above propagate unchanged).
typer.echo(
f"Error: --password could not update the Studio auth database ({exc}); not starting.",
err = True,
)
raise typer.Exit(1)
finally:
conn.close()
def _strip_seeded_bootstrap_password_or_exit(*, context: str) -> None:
"""Remove the seeded plaintext bootstrap password before a public re-exec.
Version-independent protection: a re-exec'd child of ANY version (including an
old studio-venv predating the pre-bind gate) then reads None instead of
injecting the default credential into the public page. must_change_password
stays set, so the login page still forces a change and the timer still arms.
Removal IS the protection, so if it fails (locked file, read-only auth dir)
fail closed rather than publish it.
"""
bootstrap_file = STUDIO_HOME / "auth" / BOOTSTRAP_PASSWORD_FILE
try:
bootstrap_file.unlink(missing_ok = True)
except OSError as exc:
typer.echo(
"Error: refusing to publish Studio on a public Cloudflare URL: "
f"could not remove the seeded bootstrap password file ({exc}), so an "
f"older Studio child could still serve the default credential ({context}). "
"Delete it manually or change the admin password (run `unsloth studio` "
"locally with a terminal attached, or `unsloth studio reset-password`), "
"then retry.",
err = True,
)
raise typer.Exit(1)
def _require_servable_frontend_or_exit(
*, frontend: Optional[Path], api_only: bool, cloudflare: Optional[bool], host: str, secure: bool
) -> Optional[Path]:
"""Fail closed BEFORE the pre-exposure gate if a public UI launch has no
login page to change the seeded password.
The gate strips the seeded .bootstrap_password on a headless public launch,
so if the child then cannot serve the login page the admin is locked out
(must_change_password=1, no file, no UI) until `unsloth studio reset-password`.
The login page is the ONLY in-band way to change the seeded password, so a
public non-api-only launch must have a servable dist before the strip.
Returns the dist to serve: a user-supplied --frontend (validated to contain
index.html) or the auto-resolved built dist. Returns `frontend` unchanged for
non-public or --api-only launches (no login page needed).
"""
if api_only or not _should_prompt_password_change(
cloudflare = cloudflare, host = host, secure = secure, api_only = api_only
):
return frontend
if frontend is not None:
# A user-supplied dist is not vetted by _find_frontend_dist, so verify it
# can serve the login page; else `--frontend /bad/path` bypasses the guard.
if (Path(frontend) / "index.html").is_file():
return frontend
typer.echo(
"Error: --frontend points at a directory with no index.html, so a "
"public Studio launch would have no login page to change the seeded "
"admin password. Point --frontend at a built dist, rebuild it (re-run "
"install.sh), or use --api-only.",
err = True,
)
raise typer.Exit(1)
# _find_frontend_dist only returns a path that already contains index.html.
resolved = _find_frontend_dist()
if resolved is not None:
return resolved
typer.echo(
"Error: the Studio frontend is not built, so a public launch would have "
"no login page to change the seeded admin password. Build it (re-run "
"install.sh), pass --frontend PATH to a built dist, or use --api-only.",
err = True,
)
raise typer.Exit(1)
def _validate_inproc_backend_before_strip(
*, cloudflare: Optional[bool], host: str, secure: bool, api_only: bool
) -> None:
"""In-venv (in-process) analogue of the re-exec launcher check.
In-venv there is no re-exec, so the backend is imported in-process only AFTER
the gate. On the headless public path the gate strips the seeded
.bootstrap_password, so a broken venv that fails at import would leave
must_change_password=1 with no password to log in. Import the backend up front
on that path and exit cleanly if broken, before anything is stripped.
Headless-only so an interactive prompt is not delayed behind the import.
"""
if not _should_prompt_password_change(
cloudflare = cloudflare, host = host, secure = secure, api_only = api_only
):
return
if _prompt_streams_interactive():
return
try:
_load_run_module()
except Exception as exc:
typer.echo(
f"Error: the Studio backend could not be loaded ({exc}); refusing to "
"expose Studio publicly before it is confirmed runnable. Re-run: "
"unsloth studio setup",
err = True,
)
raise typer.Exit(1)
def _tunnel_binary_confirmed_unavailable() -> bool:
"""True only if cloudflared is provably unavailable (found nowhere on PATH or
in the Studio cache AND the download failed), so the tunnel cannot start.
Used on the --secure path (loopback bind, so the tunnel is the ONLY public
exposure) to skip stripping the seeded recovery password before a public URL
that will never come up. Loads the stdlib-only cloudflare_tunnel helper by
file path so the check runs in the parent, before the strip.
Returns False on ANY uncertainty: a possible credential leak outweighs a
recoverable lockout, so the caller keeps the strip unless the tunnel is
provably dead.
"""
run_py = _find_run_py()
if run_py is None:
return False
backend_dir = run_py.parent
tunnel_py = backend_dir / "cloudflare_tunnel.py"
if not tunnel_py.is_file():
return False
# ensure_cloudflared() lazily imports utils.paths.storage_roots to resolve the
# Studio bin cache. The outer CLI hasn't added studio/backend to sys.path yet,
# so that import would fail and return None (a false "unavailable" that wrongly
# refuses --secure). Add the backend dir so the cache path resolves as in the child.
added_backend_path = False
try:
if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
added_backend_path = True
spec = importlib.util.spec_from_file_location("studio.backend.cloudflare_tunnel", tunnel_py)
if spec is None or spec.loader is None:
return False
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module.ensure_cloudflared() is None
except Exception:
return False
finally:
if added_backend_path:
try:
sys.path.remove(str(backend_dir))
except ValueError:
pass
def _child_self_suppresses(*, in_studio_venv: bool, child_run_py: Optional[Path]) -> bool:
"""True when the child that will serve Studio is provably THIS install's
backend, whose pre-bind gate sets app.state.suppress_bootstrap_injection and
so never serves the seeded credential publicly -- even with .bootstrap_password
on disk. The parent-side strip is then unnecessary and can be skipped to avoid
a lockout if the tunnel never comes up, keeping the file for LOCAL recovery.
True iff we run in-process here, or the re-exec target is the outer install's
own run.py (identity match). False on ANY doubt -- a studio-venv console script
or a venv run.py that may predate the gate -- so the strip stays in force
wherever an old child is possible.
"""
if in_studio_venv:
return True
if child_run_py is None:
return False
try:
outer_run_py = (_PACKAGE_ROOT / "studio" / "backend" / "run.py").resolve()
return child_run_py.resolve() == outer_run_py
except OSError:
return False
def _enforce_password_change_before_exposure(
*,
cloudflare: Optional[bool],
host: str,
secure: bool,
api_only: bool,
child_self_suppresses: bool = False,
) -> None:
"""Force a terminal password change before the first public (tunnel) exposure.
When the launch will start the tunnel and the admin still has its
auto-generated bootstrap password, ask for a new one in the terminal (masked,
confirmed) before any server or tunnel exists. Committing here, in the parent,
keeps the password off argv/env and an older studio-venv child sees it
immediately. Without a terminal, warn and fall back to the bootstrap shutdown
timer (~1h, UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT).
"""
if not _should_prompt_password_change(
cloudflare = cloudflare, host = host, secure = secure, api_only = api_only
):
return
# Before public exposure we must PROVE the admin password is no longer the
# seeded default. If we cannot (auth DB won't open, or a fresh admin cannot be
# seeded + committed below), an old studio-venv child could regenerate a fresh
# bootstrap credential and serve it; stripping a file we can't vouch for cannot
# stop a regeneration. So those cases fail closed, as does a failure after the
# user typed a new password.
try:
conn = _connect_auth_db()
except (OSError, sqlite3.Error) as exc:
# Cannot open the auth DB, so cannot confirm a committed admin exists.
# Refuse rather than risk a child serving the default login; a transient
# lock clears on retry.
typer.echo(
"Error: refusing to publish Studio on a public Cloudflare URL: could "
f"not open the Studio auth database ({exc}) to confirm the admin "
"password was changed. Retry (a transient database lock clears), or "
"change the password first (run `unsloth studio` locally with a "
"terminal attached, or `unsloth studio reset-password`).",
err = True,
)
raise typer.Exit(1)
try:
try:
_ensure_cli_default_admin(conn)
# Persist a freshly seeded admin before we might re-exec: the INSERT is
# otherwise uncommitted and rolls back on conn.close(). If the seed or
# commit fails, no admin is committed, so a re-exec'd OLD child finds
# none, regenerates a fresh bootstrap password + file, and serves THAT
# -- stripping cannot stop a regeneration. Can't prove a committed
# admin, so fail closed.
conn.commit()
except (OSError, sqlite3.Error) as exc:
# Best-effort remove any half-written seed file (its row rolled back);
# the launch is refused regardless.
try:
(STUDIO_HOME / "auth" / BOOTSTRAP_PASSWORD_FILE).unlink(missing_ok = True)
except OSError:
pass
typer.echo(
"Error: refusing to publish Studio on a public Cloudflare URL: could "
f"not initialize the admin account ({exc}), so a re-exec'd Studio "
"child could regenerate and serve a default credential. Retry (a "
"transient database lock clears), or change the password first (run "
"`unsloth studio` locally with a terminal attached, or `unsloth "
"studio reset-password`).",
err = True,
)
raise typer.Exit(1)
try:
row = conn.execute(
"SELECT password_salt, password_hash, must_change_password "
"FROM auth_user WHERE username = ?",
(DEFAULT_ADMIN_USERNAME,),
).fetchone()
except (OSError, sqlite3.Error) as exc:
if child_self_suppresses:
# Could not read must_change back, but the child is this install's
# own backend and suppresses the injection, so nothing serves the
# seeded credential; proceed without stripping.
return
# The admin is committed above, so an old child finds it and won't
# regenerate; we just couldn't read must_change back. Strip the seeded
# file so nothing serves it, failing closed if the strip itself fails.
typer.echo(
f"Warning: could not read the Studio admin state back ({exc}); "
"removing the seeded bootstrap password before public exposure.",
err = True,
)
_strip_seeded_bootstrap_password_or_exit(context = "auth DB row unreadable")
return
if not row or not row[2]:
return
if not _prompt_streams_interactive():
# Only proceed headless if the bootstrap shutdown deadline will protect
# the launch: it never arms for api-only, and TIMEOUT=0 disables it.
if api_only or not _bootstrap_deadline_active():
typer.echo(
"Error: 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.",
err = True,
)
raise typer.Exit(1)
if child_self_suppresses:
# The child is this install's own backend, whose pre-bind gate sets
# app.state.suppress_bootstrap_injection, so the seeded credential
# is never served publicly even with the file on disk. Skip the
# strip: unnecessary here, and it would lock the user out if the
# tunnel never comes up (e.g. a --secure loopback whose tunnel
# fails). Keep the file for LOCAL recovery; must_change stays set
# and the deadline arms.
typer.echo(
"Warning: Studio is being exposed publicly while the admin "
"account still uses its auto-generated bootstrap password. The "
"login page forces a change and the credential is never served "
"on the public page. Set a new password by running `unsloth "
"studio` locally with a terminal attached, or `unsloth studio "
"reset-password`; Studio shuts down after ~1h if the password "
"stays unchanged (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT).",
err = True,
)
return
# The strip permanently removes the only plaintext recovery credential.
# On --secure the bind is loopback, so the tunnel is the ONLY public
# exposure: if cloudflared is provably unavailable no public URL can
# start, so stripping would just lock the user out. Refuse with the
# credential preserved. (A wildcard --cloudflare bind is public
# regardless of the tunnel, so it still strips below, as does any
# uncertainty.)
if secure and _tunnel_binary_confirmed_unavailable():
typer.echo(
"Error: refusing to expose Studio: the Cloudflare tunnel binary "
"(cloudflared) is unavailable and could not be downloaded, so no "
"public URL can start. The seeded bootstrap password is preserved "
"for recovery; fix connectivity and retry, or change the password "
"first (`unsloth studio` locally, or `unsloth studio "
"reset-password`).",
err = True,
)
raise typer.Exit(1)
# Mixed-version safety: an OLD studio-venv child (predating this gate)
# has no pre-bind suppression and would read the seeded credential back
# from disk and inject it into the public HTML until the deadline.
# Delete the file here, in the parent, so a fresh child of ANY version
# reads None. must_change_password stays set, so the login page still
# forces a change and the timer still arms; only the on-disk copy goes.
_strip_seeded_bootstrap_password_or_exit(context = "no terminal to change it")
typer.echo(
"Warning: Studio is being exposed publicly while the admin account "
"still uses its auto-generated bootstrap password. The seeded password "
"file has been removed so it is not served on the public page. Set a new "
"password by running `unsloth studio` locally with a terminal attached, "
"or `unsloth studio reset-password`; Studio shuts down after ~1h if the "
"password stays unchanged (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT).",
err = True,
)
return
password_salt, password_hash = row[0], row[1]
def _is_current_password(candidate: str) -> bool:
return hmac.compare_digest(
_pbkdf2_hex(candidate, password_salt.encode("utf-8")), password_hash
)
typer.echo(
"Unsloth Studio will be exposed on the public internet, so set a "
"password now. Ctrl+C to abort.",
err = True,
)
try:
new_password = _password_prompt.prompt_new_password(_is_current_password)
except (KeyboardInterrupt, EOFError):
typer.echo(
"\nError: password change aborted; refusing to expose Studio "
"with the default admin password. Re-run and set a password, "
"or launch without --secure/--cloudflare.",
err = True,
)
raise typer.Exit(1)
_cli_update_password(conn, DEFAULT_ADMIN_USERNAME, new_password)
typer.echo(f"Password updated for '{DEFAULT_ADMIN_USERNAME}'.", err = True)
finally:
conn.close()
def _load_model_via_http(
port: int,
api_key: str,
@ -713,13 +1244,13 @@ def studio_default(
f"defaults to {_PARALLEL_DEFAULT_RUN}."
),
),
cloudflare: bool = typer.Option(
True,
cloudflare: Optional[bool] = typer.Option(
None,
"--cloudflare/--no-cloudflare",
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.",
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 forces "
"it off but does not change a raw wildcard bind.",
),
secure: bool = typer.Option(
False,
@ -747,6 +1278,14 @@ def studio_default(
help = "Force server-side tools (web search, code execution) on or off for "
"every request. Default: on for every bind, with the per-chat UI toggle honored.",
),
password: str = typer.Option(
"",
"--password",
help = "Set the INITIAL admin password non-interactively (headless setups), "
"only when none is set yet. Also reads the UNSLOTH_STUDIO_PASSWORD env var, or "
"`--password -` to read one line from stdin. A literal value is visible in the "
"process list and shell history. Rotate later with `unsloth studio reset-password`.",
),
):
"""Launch the Unsloth Studio server."""
# Back-compat: --not-secure is a deprecated alias for --no-secure.
@ -766,13 +1305,14 @@ def studio_default(
err = True,
)
raise typer.Exit(2)
# Same for --no-cloudflare: it would not reach the subcommand.
if not cloudflare:
# Same for --cloudflare/--no-cloudflare: it would not reach the subcommand.
if cloudflare is not None:
_cf_flag = "--cloudflare" if cloudflare else "--no-cloudflare"
typer.echo(
f"Error: --no-cloudflare on `unsloth studio` applies to the "
f"Error: {_cf_flag} on `unsloth studio` applies to the "
f"plain-server path only. For `unsloth studio "
f"{ctx.invoked_subcommand}`, put it after the subcommand: "
f"`unsloth studio {ctx.invoked_subcommand} --no-cloudflare ...`",
f"`unsloth studio {ctx.invoked_subcommand} {_cf_flag} ...`",
err = True,
)
raise typer.Exit(2)
@ -817,17 +1357,34 @@ def studio_default(
err = True,
)
raise typer.Exit(2)
# Same for --password: it applies to the plain-server path only.
if password:
typer.echo(
f"Error: --password on `unsloth studio` applies to the "
f"plain-server path only. For `unsloth studio "
f"{ctx.invoked_subcommand}`, put it after the subcommand: "
f"`unsloth studio {ctx.invoked_subcommand} --password ...`",
err = True,
)
raise typer.Exit(2)
return
# --secure requires the tunnel; force a loopback bind.
if secure:
if not cloudflare:
if cloudflare is False:
typer.echo(
"Error: --secure requires the Cloudflare tunnel; do not combine it "
"with --no-cloudflare.",
err = True,
)
raise typer.Exit(2)
if host not in ("127.0.0.1", "localhost", "::1"):
typer.echo(
"Note: --secure ignores -H (it binds loopback and serves only "
"through the Cloudflare tunnel). Drop --secure to bind "
f"{host} directly, or keep --secure for a tunnel-only public link.",
err = True,
)
host = "127.0.0.1"
# --verbose restores the per-request access logs that are suppressed by
@ -835,13 +1392,76 @@ def studio_default(
if verbose:
_enable_verbose_access_logs()
# Use the studio venv if it exists and we aren't already in it.
# Use the studio venv if present and not already in it. Resolve the child
# launcher BEFORE the gate: a headless gate strips the seeded
# .bootstrap_password, so aborting afterward (venv/run.py missing) would leave
# must_change_password=1 with no password to log in.
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
studio_python = run_py = None
resolved_frontend = frontend
if not in_studio_venv:
studio_python = _studio_venv_python()
run_py = _find_run_py()
if not (studio_python and run_py):
typer.echo("Studio not set up. Run install.sh first.")
raise typer.Exit(1)
# A public UI launch must have a servable login page BEFORE the gate can
# strip the seeded .bootstrap_password, or the child has no way to change
# it. Also returns the resolved dist so the child serves a real build
# regardless of where its __file__ lands (fixes the shadowed silent 404).
resolved_frontend = _require_servable_frontend_or_exit(
frontend = resolved_frontend,
api_only = api_only,
cloudflare = cloudflare,
host = host,
secure = secure,
)
# Non-public / api-only launches skip that validation but still forward an
# explicitly resolved dist for the same silent-404 reason.
if resolved_frontend is None and not api_only:
resolved_frontend = _find_frontend_dist()
else:
# Already in the studio venv: no re-exec, served in-process below. On the
# headless public path the gate strips the seeded .bootstrap_password, so
# validate BOTH FIRST -- else a bad dist or broken venv fails only after
# the strip (must_change_password=1, no password to log in). Frontend check
# first (cheap); the backend import is headless-only so an interactive
# prompt is not delayed behind it.
resolved_frontend = _require_servable_frontend_or_exit(
frontend = resolved_frontend,
api_only = api_only,
cloudflare = cloudflare,
host = host,
secure = secure,
)
_validate_inproc_backend_before_strip(
cloudflare = cloudflare, host = host, secure = secure, api_only = api_only
)
# A supplied --password / UNSLOTH_STUDIO_PASSWORD / stdin sets the initial
# admin password here in the parent, before the gate and any re-exec, so the
# secret never reaches the child argv; strip the env var so a re-exec'd child
# can't re-read it. The interactive gate below then no-ops.
_apply_supplied_password_before_launch(_password_prompt.resolve_supplied_password(password))
os.environ.pop(_password_prompt.SUPPLIED_PASSWORD_ENV, None)
# Public (tunnel) exposure with the seeded default password: force a terminal
# password change first, before any re-exec or server exists. The child is
# self-suppressing when we serve in-process or re-exec this install's own
# run.py (its pre-bind gate suppresses the injection), so the gate can skip
# the destructive strip.
_enforce_password_change_before_exposure(
cloudflare = cloudflare,
host = host,
secure = secure,
api_only = api_only,
child_self_suppresses = _child_self_suppresses(
in_studio_venv = in_studio_venv, child_run_py = run_py
),
)
if not in_studio_venv:
if studio_python and run_py:
if not silent:
typer.echo("Launching Unsloth Studio... Please wait...")
@ -855,20 +1475,22 @@ def studio_default(
"--parallel",
str(parallel),
]
# Resolve frontend explicitly so the spawned run.py uses a real
# built dist regardless of where its __file__ lands. Skip in
# --api-only (no UI served).
resolved_frontend = frontend
if resolved_frontend is None and not api_only:
resolved_frontend = _find_frontend_dist()
# Forward the frontend dist resolved before the gate (skipped in
# --api-only, which serves no UI).
if resolved_frontend is not None:
args.extend(["--frontend", str(resolved_frontend)])
if silent:
args.append("--silent")
if api_only:
args.append("--api-only")
# Forward the explicit polarity (matches run.py's BooleanOptionalAction).
args.append("--cloudflare" if cloudflare else "--no-cloudflare")
# Forward polarity explicitly: _find_run_py can fall back to an older
# run.py (--cloudflare defaulted on), so an unset default must not let a
# mixed install silently re-enable the tunnel. --secure implies it, so
# forward nothing then.
if cloudflare is True:
args.append("--cloudflare")
elif not secure:
args.append("--no-cloudflare")
args.append("--secure" if secure else "--no-secure")
# Forward an explicit tool policy; None -> run.py leaves it unset (tools on).
if enable_tools is True:
@ -920,8 +1542,10 @@ def studio_default(
secure = secure,
enable_tools = enable_tools,
)
if frontend is not None:
run_kwargs["frontend_path"] = frontend
# Forward the frontend validated before the gate (in-venv path), so the
# in-process server serves exactly the dist we vouched for.
if resolved_frontend is not None:
run_kwargs["frontend_path"] = resolved_frontend
run_server(**run_kwargs)
try:
@ -1106,13 +1730,13 @@ def run(
f"{_PARALLEL_DEFAULT_RUN} (pre-PR hardcoded value)."
),
),
cloudflare: bool = typer.Option(
True,
cloudflare: Optional[bool] = typer.Option(
None,
"--cloudflare/--no-cloudflare",
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.",
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 forces "
"it off but does not change a raw wildcard bind.",
),
secure: bool = typer.Option(
False,
@ -1136,6 +1760,14 @@ def run(
"decode speed, MoE usually don't."
),
),
password: str = typer.Option(
"",
"--password",
help = "Set the INITIAL admin password non-interactively (headless setups), "
"only when none is set yet. Also reads the UNSLOTH_STUDIO_PASSWORD env var, or "
"`--password -` to read one line from stdin. A literal value is visible in the "
"process list and shell history. Rotate later with `unsloth studio reset-password`.",
),
):
"""Start Studio, load a model, print an API key -- one-liner server.
@ -1207,13 +1839,20 @@ def run(
# --secure requires the tunnel; force a loopback bind so the raw port is never public.
if secure:
if not cloudflare:
if cloudflare is False:
typer.echo(
"Error: --secure requires the Cloudflare tunnel; do not combine it "
"with --no-cloudflare.",
err = True,
)
raise typer.Exit(2)
if host not in ("127.0.0.1", "localhost", "::1"):
typer.echo(
"Note: --secure ignores -H (it binds loopback and serves only "
"through the Cloudflare tunnel). Drop --secure to bind "
f"{host} directly, or keep --secure for a tunnel-only public link.",
err = True,
)
host = "127.0.0.1"
# Tool policy no longer depends on the bind: tools default on everywhere
@ -1228,10 +1867,14 @@ def run(
silent = silent,
)
# 1. Re-exec into the studio venv (same pattern as studio_default).
# 1. Re-exec into the studio venv (same pattern as studio_default). Resolve
# the child launcher BEFORE the gate: a headless gate strips the seeded
# .bootstrap_password, so aborting afterward (venv/entry point missing) would
# leave must_change_password=1 with no password to log in.
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
studio_bin = None
resolved_frontend = frontend
if not in_studio_venv:
studio_python = _studio_venv_python()
if not studio_python:
@ -1242,6 +1885,56 @@ def run(
if not studio_bin.is_file():
typer.echo("Studio venv missing 'unsloth' entry point. Re-run: unsloth studio setup")
raise typer.Exit(1)
# `run` serves the same Studio UI (unless --api-only); a public launch must
# have a servable login page BEFORE the gate strips the seeded password, or
# the child has no way to change it. Validate here and forward the resolved
# dist so a shadowed child that can't self-resolve one still serves it.
resolved_frontend = _require_servable_frontend_or_exit(
frontend = frontend,
api_only = api_only,
cloudflare = cloudflare,
host = host,
secure = secure,
)
else:
# In-venv (in-process) run: validate the servable frontend and importable
# backend before the headless gate strips the seeded password. Frontend
# check first (cheap); backend import is headless-only so a prompt isn't
# delayed.
resolved_frontend = _require_servable_frontend_or_exit(
frontend = frontend,
api_only = api_only,
cloudflare = cloudflare,
host = host,
secure = secure,
)
_validate_inproc_backend_before_strip(
cloudflare = cloudflare, host = host, secure = secure, api_only = api_only
)
# A supplied --password / UNSLOTH_STUDIO_PASSWORD / stdin sets the initial
# admin password here in the parent, before the gate and any re-exec, so the
# secret never reaches the child argv; strip the env var so a re-exec'd child
# can't re-read it. The interactive gate below then no-ops.
_apply_supplied_password_before_launch(_password_prompt.resolve_supplied_password(password))
os.environ.pop(_password_prompt.SUPPLIED_PASSWORD_ENV, None)
# Public (tunnel) exposure with the seeded default password: force a terminal
# password change first, before any re-exec or server exists. The re-exec here
# runs the studio venv's `unsloth` console script (a possibly-OLD child), so it
# is NOT provably self-suppressing -- only the in-process case is, and the
# strip stays in force otherwise.
_enforce_password_change_before_exposure(
cloudflare = cloudflare,
host = host,
secure = secure,
api_only = api_only,
child_self_suppresses = _child_self_suppresses(
in_studio_venv = in_studio_venv, child_run_py = None
),
)
if not in_studio_venv:
args = [
str(studio_bin),
"studio",
@ -1262,8 +1955,12 @@ def run(
# Forward the explicit polarity; a future default flip on one
# layer must not silently invert behaviour for the other.
args.append("--load-in-4bit" if load_in_4bit else "--no-load-in-4bit")
if frontend:
args.extend(["--frontend", str(frontend)])
# Forward the frontend resolved before the gate, not just a user-supplied
# one: the parent may have found a built dist the shadowed child cannot,
# and stripping without forwarding it would abort the child at frontend
# setup (lockout).
if resolved_frontend is not None:
args.extend(["--frontend", str(resolved_frontend)])
if api_only:
args.append("--api-only")
if silent:
@ -1279,8 +1976,13 @@ def run(
# Typer claims --parallel outside ctx.args; without this the
# child reverts to its default and silently drops the value.
args.extend(["--parallel", str(parallel)])
# Forward the explicit polarity (same rationale as --load-in-4bit above).
args.append("--cloudflare" if cloudflare else "--no-cloudflare")
# Always forward explicit polarity: a mixed-version studio venv whose old
# default was --cloudflare-on must not silently re-enable the tunnel.
# --secure implies it, so forward nothing then.
if cloudflare is True:
args.append("--cloudflare")
elif not secure:
args.append("--no-cloudflare")
args.append("--secure" if secure else "--no-secure")
args.append("--tensor-parallel" if tensor_parallel else "--no-tensor-parallel")
if verbose:
@ -1322,8 +2024,9 @@ def run(
# TAURI_PORT line would corrupt that machine-parseable output.
emit_tauri_port = False,
)
if frontend is not None:
run_kwargs["frontend_path"] = frontend
# Forward the frontend validated before the gate (in-venv path).
if resolved_frontend is not None:
run_kwargs["frontend_path"] = resolved_frontend
app = run_server(**run_kwargs)
actual_port = getattr(app.state, "server_port", port) or port
@ -1943,9 +2646,44 @@ def reset_password():
]
had_db = db_file.exists()
db_file.unlink(missing_ok = True)
# Delete auth.db FIRST and prove it is gone before touching the seeded
# credential files. If it cannot be removed (a running Studio or Windows
# holds it open, or a read-only auth dir), abort with the credential files
# untouched: deleting them while an un-resettable DB (must_change_password=1)
# survives would lock a forgotten-password reset out of any recovery
# credential. Failing here leaves a consistent, still-recoverable state.
try:
db_file.unlink(missing_ok = True)
except OSError as exc:
typer.echo(
f"Error: could not delete the auth database ({exc}). Stop any running "
"Studio and retry; no credential files were changed.",
err = True,
)
raise typer.Exit(1)
# The DB is gone, so the next start re-seeds. Invalidate the seeded plaintext
# credential files so that re-seed generates a FRESH password instead of
# reusing a stale one: unlink only ignores FileNotFoundError, so a
# locked/undeletable file (Windows AV, read-only dir) would otherwise survive
# and generate_bootstrap_password() would read it back and re-validate the
# credential this reset revoked. Truncate on unlink failure; if a file can be
# neither removed nor truncated, fail closed -- the DB is already gone, so a
# surviving plaintext would be reused, and the user must remove it manually.
for path in stale_files:
path.unlink(missing_ok = True)
try:
path.unlink(missing_ok = True)
except OSError:
try:
path.write_text("")
except OSError as exc:
typer.echo(
f"Error: could not remove or clear {path.name} ({exc}); delete "
"it manually before restarting Studio or the old password may "
"be reused.",
err = True,
)
raise typer.Exit(1)
if not had_db:
typer.echo("No auth database found -- nothing to reset.")

View file

@ -294,17 +294,58 @@ def test_merge_codex_config_keeps_user_oss_provider():
assert _parse_toml(merged)["oss_provider"] == "ollama"
def test_write_codex_config_profile(tmp_path):
def test_write_codex_config_profile(tmp_path, monkeypatch):
monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True)
start.write_codex_config(BASE, MODEL, tmp_path)
profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text())
assert profile["oss_provider"] == "unsloth_api"
assert profile["model_provider"] == "unsloth_api"
assert profile["model"] == MODEL["id"]
assert profile["model_context_window"] == 131072
catalog_path = Path(profile["model_catalog_json"])
assert catalog_path == Path("model-catalog.json")
catalog = json.loads((tmp_path / catalog_path).read_text())
assert catalog["models"][0]["slug"] == MODEL["id"]
assert catalog["models"][0]["context_window"] == 131072
assert catalog["models"][0]["max_context_window"] == 131072
assert catalog["models"][0]["supports_reasoning_summary_parameter"] is False
assert catalog["models"][0]["supports_parallel_tool_calls"] is False
assert catalog["models"][0]["base_instructions"] == start._CODEX_FALLBACK_PROMPT.read_text()
config = _parse_toml((tmp_path / "config.toml").read_text())
assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN"
def test_write_codex_config_catalog_without_context_length(tmp_path, monkeypatch):
monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True)
start.write_codex_config(BASE, {"id": "unsloth/no-window"}, tmp_path)
profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text())
catalog = json.loads((tmp_path / profile["model_catalog_json"]).read_text())
entry = catalog["models"][0]
assert entry["slug"] == "unsloth/no-window"
assert "context_window" not in entry
assert "max_context_window" not in entry
@pytest.mark.parametrize(
("version", "expected"),
[("codex-cli 0.109.0", False), ("codex-cli 0.110.0", True), ("codex-cli 0.144.4", True)],
)
def test_codex_model_catalog_version_gate(monkeypatch, version, expected):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: version)
assert start._codex_supports_model_catalog() is expected
def test_write_codex_config_omits_catalog_for_old_codex(tmp_path, monkeypatch):
monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: False)
start.write_codex_config(BASE, MODEL, tmp_path)
profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text())
assert "model_catalog_json" not in profile
assert not (tmp_path / "model-catalog.json").exists()
@pytest.fixture()
def fake_studio(tmp_path, monkeypatch):
calls = []
@ -742,7 +783,12 @@ def test_opencode_inline_config_beats_project_config(fake_studio):
assert result.exit_code == 0, result.output
inline = _opencode_inline_config(result.output)
assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"
assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"}
assert inline["permission"] == {
"edit": "allow",
"bash": "allow",
"webfetch": "allow",
"external_directory": {"*": "allow"},
}
assert "sk-unsloth" not in result.output # key stays in the private file, not the env
@ -1611,12 +1657,50 @@ def test_write_openclaw_config_fresh(tmp_path):
]
# The default model must be pinned or OpenClaw has nothing active.
assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}"
assert config["agents"]["defaults"]["workspace"] == str(tmp_path / "workspace")
assert (tmp_path / "workspace").is_dir()
assert config["gateway"]["mode"] == "local"
assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway
if os.name != "nt": # the file holds an API key
assert path.stat().st_mode & 0o777 == 0o600
def test_write_openclaw_config_clears_per_agent_path_overrides(tmp_path):
path = tmp_path / "openclaw.json"
path.write_text(
json.dumps(
{
"agents": {
"defaults": {"workspace": "/old/default"},
"list": [
{
"id": "main",
"default": True,
"workspace": "/old/main-workspace",
"agentDir": "/old/main-agent",
"model": "keep/me",
},
{
"id": "reviewer",
"workspace": "/old/reviewer-workspace",
"agentDir": "/old/reviewer-agent",
},
],
}
}
)
)
start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path)
agents = json.loads(path.read_text())["agents"]
assert agents["defaults"]["workspace"] == str(tmp_path / "workspace")
assert agents["list"] == [
{"id": "main", "default": True, "model": "keep/me"},
{"id": "reviewer"},
]
def test_write_openclaw_config_preserves_and_idempotent(tmp_path):
path = tmp_path / "openclaw.json"
path.write_text(
@ -1660,11 +1744,32 @@ def test_connect_openclaw_no_launch(fake_studio, tmp_path):
config = json.loads(config_path.read_text())
assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface"
assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}"
assert config["agents"]["defaults"]["workspace"] == str(
tmp_path / "agents" / "openclaw" / "workspace"
)
assert _launch_command(result.output) == ["openclaw", "tui", "--local"]
# OpenAI /v1/chat/completions works on either backend — no GGUF gate.
assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
def test_connect_openclaw_wsl_windows_shim_translates_workspace(fake_studio, tmp_path, monkeypatch):
windows_workspace = r"\\wsl.localhost\Ubuntu\tmp\openclaw\workspace"
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
monkeypatch.setattr(
start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/openclaw"
)
monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: windows_workspace)
result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"])
assert result.exit_code == 0, result.output
config_path = tmp_path / "agents" / "openclaw" / "openclaw.json"
config = json.loads(config_path.read_text())
assert config["agents"]["defaults"]["workspace"] == windows_workspace
assert (config_path.parent / "workspace").is_dir()
def test_connect_openclaw_no_launch_keeps_explicit_subcommand(fake_studio):
result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch", "crestodian"])
assert result.exit_code == 0, result.output
@ -2090,7 +2195,12 @@ def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path):
result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"])
assert result.exit_code == 0, result.output
config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text())
assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"}
assert config["permission"] == {
"edit": "allow",
"bash": "allow",
"webfetch": "allow",
"external_directory": {"*": "allow"},
}
def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path):
@ -2112,6 +2222,7 @@ def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path):
"edit": "allow",
"bash": "allow",
"webfetch": "allow",
"external_directory": {"*": "allow"},
}
plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert plain.exit_code == 0, plain.output
@ -2119,6 +2230,7 @@ def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path):
"edit": "ask",
"bash": "ask",
"webfetch": "ask",
"external_directory": {"*": "ask"},
}
@ -2152,7 +2264,12 @@ def test_write_opencode_config_yolo_unit(tmp_path):
path = tmp_path / "opencode.json"
start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True)
config = json.loads(path.read_text())
assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"}
assert config["permission"] == {
"edit": "allow",
"bash": "allow",
"webfetch": "allow",
"external_directory": {"*": "allow"},
}
def test_write_openclaw_config_yolo_unit(tmp_path):
@ -2180,7 +2297,12 @@ def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp
config = json.loads(config_path.read_text())
# The yolo allow policy is replaced by a prompting one, not deleted (which would
# revert to OpenCode's permissive "allow" default).
assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"}
assert config["permission"] == {
"edit": "ask",
"bash": "ask",
"webfetch": "ask",
"external_directory": {"*": "ask"},
}
# The session provider survives the cleanup.
assert start._OPENCODE_PROVIDER in config["provider"]
@ -2218,7 +2340,12 @@ def test_write_opencode_config_yolo_then_plain_unit(tmp_path):
start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False)
config = json.loads(path.read_text())
# A plain rerun replaces the yolo allow policy with a prompting one.
assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"}
assert config["permission"] == {
"edit": "ask",
"bash": "ask",
"webfetch": "ask",
"external_directory": {"*": "ask"},
}
def test_openclaw_non_yolo_keeps_runtime_approvals(tmp_path):
@ -2670,8 +2797,7 @@ def test_default_launch_has_no_resume_token(fake_studio, monkeypatch):
def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch):
# openclaw/hermes persist their session dir but have no non-interactive resume
# selector, so --persist must not append a token; their own picker resumes.
# Persistence alone must not select a session.
for agent in ("openclaw", "hermes"):
monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}")
captured = _capture_launch(monkeypatch, [agent, "--persist"])
@ -2679,6 +2805,135 @@ def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatc
assert "--continue" not in captured["command"]
@pytest.mark.parametrize(
("args", "expected"),
[
(
["--resume", "session-id", "-z", "follow up"],
[
"chat",
"-Q",
"--yolo",
"--accept-hooks",
"--resume",
"session-id",
"-q",
"follow up",
],
),
(
["-rsession-id", "-zfollow up"],
["chat", "-Q", "--yolo", "--accept-hooks", "-rsession-id", "-qfollow up"],
),
(
["-c=project", "-z=follow up"],
["chat", "-Q", "--yolo", "--accept-hooks", "-c=project", "-q=follow up"],
),
(
["-r", "session-id", "--oneshot=follow up"],
[
"chat",
"-Q",
"--yolo",
"--accept-hooks",
"-r",
"session-id",
"--query=follow up",
],
),
(
["--continue", "project", "--oneshot", "follow up"],
[
"chat",
"-Q",
"--yolo",
"--accept-hooks",
"--continue",
"project",
"-q",
"follow up",
],
),
(
["--yolo", "--resume", "session-id", "-z", "follow up"],
[
"chat",
"-Q",
"--accept-hooks",
"--yolo",
"--resume",
"session-id",
"-q",
"follow up",
],
),
(
["--accept-hooks", "--resume", "session-id", "-z", "follow up"],
[
"chat",
"-Q",
"--yolo",
"--accept-hooks",
"--resume",
"session-id",
"-q",
"follow up",
],
),
(
["--resume", "chat", "-z", "follow up"],
[
"chat",
"-Q",
"--yolo",
"--accept-hooks",
"--resume",
"chat",
"-q",
"follow up",
],
),
(["--resume", "session-id"], ["--resume", "session-id"]),
(["-z", "new session"], ["-z", "new session"]),
],
)
def test_hermes_resume_oneshot_args(args, expected):
assert start._hermes_resume_oneshot_args(args) == expected
def test_hermes_resume_oneshot_uses_session_aware_chat(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/hermes")
captured = _capture_launch(
monkeypatch,
["hermes", "--persist", "--resume", "session-id", "-z", "follow up"],
)
assert captured["command"][1:] == [
"chat",
"-Q",
"--yolo",
"--accept-hooks",
"--resume",
"session-id",
"-q",
"follow up",
]
@pytest.mark.parametrize("usage_arg", ["--usage-file", "--usage-file=usage.json"])
def test_hermes_resume_oneshot_rejects_usage_file(monkeypatch, usage_arg):
monkeypatch.setattr(
start,
"_connect",
lambda *args, **kwargs: pytest.fail("argument validation must run before connect"),
)
argv = ["hermes", "--resume", "session-id", "-z", "follow up", usage_arg]
if usage_arg == "--usage-file":
argv.append("usage.json")
result = CliRunner().invoke(start.start_app, argv)
assert result.exit_code == 2
assert "cannot resume a one-shot session with --usage-file" in result.output
def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
# The persistence flag is --persist, NOT --resume, so an agent's own
# `--resume <id>` (e.g. `unsloth start claude --resume <guid>`) still flows

View file

@ -3,8 +3,8 @@
"""Tests for the `--cloudflare/--no-cloudflare` Studio flag.
Pins the typer Option (default on) on both `unsloth studio` and
`unsloth studio run`, and that the chosen polarity reaches the re-exec'd
Pins the typer Option (tri-state, default off / None) on both `unsloth studio`
and `unsloth studio run`, and that the chosen polarity reaches the re-exec'd
child and run_server. Modeled on test_studio_run_parallel_flag.py.
"""
@ -33,7 +33,7 @@ _BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"]
# ── option registration ──────────────────────────────────────────────
def test_run_exposes_cloudflare_option_default_on():
def test_run_exposes_cloudflare_option_default_off():
import inspect
sig = inspect.signature(_studio().run)
@ -41,16 +41,16 @@ def test_run_exposes_cloudflare_option_default_on():
opt = sig.parameters["cloudflare"].default
decls = set(getattr(opt, "param_decls", []) or [])
assert "--cloudflare/--no-cloudflare" in decls
assert getattr(opt, "default", None) is True
assert getattr(opt, "default", "missing") is None
def test_studio_default_exposes_cloudflare_option_default_on():
def test_studio_default_exposes_cloudflare_option_default_off():
import inspect
sig = inspect.signature(_studio().studio_default)
assert "cloudflare" in sig.parameters
opt = sig.parameters["cloudflare"].default
assert getattr(opt, "default", None) is True
assert getattr(opt, "default", "missing") is None
# ── re-exec forwarding: `unsloth studio run` ─────────────────────────
@ -69,6 +69,11 @@ def _install_run_reexec_capture(monkeypatch, *, platform = "linux"):
monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
fake_venv = Path("/fake/studio/venv/unsloth_studio")
monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
# A built frontend dist is present so the public-launch UI check passes
# deterministically (independent of whether the repo dist was built).
monkeypatch.setattr(
studio_mod, "_find_frontend_dist", lambda: Path("/fake/studio/frontend/dist")
)
fake_bin = fake_venv / "bin" / "unsloth"
real_is_file = Path.is_file
monkeypatch.setattr(
@ -107,19 +112,23 @@ def _invoke_run(monkeypatch, args):
@pytest.mark.parametrize(
"user_flag,expected,unexpected",
"extra_flags,expected,unexpected",
[
(None, "--cloudflare", "--no-cloudflare"), # default on
("--cloudflare", "--cloudflare", "--no-cloudflare"),
("--no-cloudflare", "--no-cloudflare", "--cloudflare"),
# Default (no flag) forwards --no-cloudflare explicitly so a mixed-version
# child venv (old default: --cloudflare on) can't re-enable the tunnel.
([], "--no-cloudflare", "--cloudflare"),
(["--cloudflare"], "--cloudflare", "--no-cloudflare"),
(["--no-cloudflare"], "--no-cloudflare", "--cloudflare"),
# --secure implies the tunnel; never forward --no-cloudflare with it.
(["--secure"], None, "--no-cloudflare"),
],
)
def test_run_reexec_forwards_cloudflare_polarity(monkeypatch, user_flag, expected, unexpected):
extras = [user_flag] if user_flag else []
captured = _invoke_run(monkeypatch, _BASE + extras)
def test_run_reexec_forwards_cloudflare_polarity(monkeypatch, extra_flags, expected, unexpected):
captured = _invoke_run(monkeypatch, _BASE + extra_flags)
assert len(captured) == 1, captured
argv = captured[0]
assert expected in argv, f"expected {expected} in child argv; got {argv}"
if expected is not None:
assert expected in argv, f"expected {expected} in child argv; got {argv}"
assert unexpected not in argv, f"unexpected {unexpected} in child argv; got {argv}"
@ -142,7 +151,11 @@ def _invoke_studio_default(
fake_venv = Path("/fake/studio/venv/unsloth_studio")
monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
monkeypatch.setattr(studio_mod, "_find_run_py", lambda: Path("/fake/studio/run.py"))
monkeypatch.setattr(studio_mod, "_find_frontend_dist", lambda: None)
# A built frontend dist is present so the public-launch UI check passes; this
# suite exercises flag forwarding, not the missing-dist lockout guard.
monkeypatch.setattr(
studio_mod, "_find_frontend_dist", lambda: Path("/fake/studio/frontend/dist")
)
monkeypatch.setattr(sys, "platform", platform)
def fake_execvp(file, argv):
@ -158,18 +171,24 @@ def _invoke_studio_default(
@pytest.mark.parametrize(
"user_flag,expected,unexpected",
"extra_flags,expected,unexpected",
[
(None, "--cloudflare", "--no-cloudflare"),
("--no-cloudflare", "--no-cloudflare", "--cloudflare"),
# Default (no flag) forwards --no-cloudflare explicitly: _find_run_py can fall
# back to an older studio-venv run.py (default on), so a mixed install must
# not re-enable the tunnel.
([], "--no-cloudflare", "--cloudflare"),
(["--cloudflare"], "--cloudflare", "--no-cloudflare"),
(["--no-cloudflare"], "--no-cloudflare", "--cloudflare"),
# --secure implies the tunnel; never forward --no-cloudflare with it.
(["--secure"], None, "--no-cloudflare"),
],
)
def test_studio_default_reexec_forwards_cloudflare(monkeypatch, user_flag, expected, unexpected):
extras = [user_flag] if user_flag else []
captured = _invoke_studio_default(monkeypatch, ["-H", "0.0.0.0"] + extras)
def test_studio_default_reexec_forwards_cloudflare(monkeypatch, extra_flags, expected, unexpected):
captured = _invoke_studio_default(monkeypatch, ["-H", "0.0.0.0"] + extra_flags)
assert len(captured) == 1, captured
argv = captured[0]
assert expected in argv, f"expected {expected}; got {argv}"
if expected is not None:
assert expected in argv, f"expected {expected}; got {argv}"
assert unexpected not in argv, f"unexpected {unexpected}; got {argv}"
@ -182,7 +201,10 @@ class _RunServerCaptured(SystemExit):
self.kwargs = dict(kwargs)
@pytest.mark.parametrize("user_flag,expected", [(None, True), ("--no-cloudflare", False)])
@pytest.mark.parametrize(
"user_flag,expected",
[(None, None), ("--cloudflare", True), ("--no-cloudflare", False)],
)
def test_run_in_venv_passes_cloudflare_to_run_server(monkeypatch, user_flag, expected):
import types
@ -348,21 +370,22 @@ def test_run_silent_emits_cloudflare_notice_for_external_bind(monkeypatch):
assert ("print", {"secure": False, "loopback_host": "127.0.0.1"}) in calls
# ── parent-level --no-cloudflare with a subcommand is rejected ──────
# ── parent-level --cloudflare/--no-cloudflare with a subcommand is rejected ─
def test_studio_default_rejects_no_cloudflare_with_subcommand(monkeypatch):
# `unsloth studio --no-cloudflare run ...` would not reach the subcommand,
# so it must error (mirrors --parallel) rather than silently still tunnel.
@pytest.mark.parametrize("flag", ["--cloudflare", "--no-cloudflare"])
def test_studio_default_rejects_cloudflare_flag_with_subcommand(monkeypatch, flag):
# `unsloth studio --cloudflare run ...` (or --no-cloudflare) would not reach the
# subcommand, so it must error (mirrors --parallel) rather than silently drop it.
import typer as _typer
studio_mod = _studio()
app = _typer.Typer()
app.add_typer(studio_mod.studio_app, name = "studio")
result = CliRunner().invoke(app, ["studio", "--no-cloudflare", "run", "--model", "X"])
result = CliRunner().invoke(app, ["studio", flag, "run", "--model", "X"])
assert result.exit_code == 2, result.output
combined = (result.output or "") + (getattr(result, "stderr", "") or "")
assert "--no-cloudflare" in combined, combined
assert flag in combined, combined
# ── run() tears the server + tunnel down if startup aborts ───────────

File diff suppressed because it is too large Load diff

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