diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 1bb4c2bb58..fa84471d36 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -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 \ diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index f540c11da4..cf8c021e38 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -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)" diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 03c0a8580d..d3d765aa84 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -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}") diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 0453c9212a..233292f7a3 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -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}") diff --git a/README.md b/README.md index 5f1630e2ba..ef45b91430 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/install.ps1 b/install.ps1 index 100a3177ba..4fa01bfa28 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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 "" } } diff --git a/install.sh b/install.sh index 3bf6fd1855..f277d0cbfd 100755 --- a/install.sh +++ b/install.sh @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 2b79121c82..917247c216 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index a0da2b2096..9bb3ab5735 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -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() diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py new file mode 100644 index 0000000000..8491019ae9 --- /dev/null +++ b/studio/backend/auth/terminal_prompt.py @@ -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 diff --git a/studio/backend/colab.py b/studio/backend/colab.py index dd274399bc..e04543b3aa 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -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, diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 671ef363f5..31bbbdc748 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -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) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 7828116236..08993a9a08 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -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. diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 093a92e38d..6b6c5373eb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 6287b184a6..163ade10c4 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -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 diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 2b1ceca75a..c2082bc198 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index a18d2758ba..c1fffb71cb 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -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() diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 5f9f1bf53d..8f9241d51f 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -6,6 +6,7 @@ import ast import codecs +import fnmatch import http.client import os import signal @@ -152,6 +153,42 @@ _COMMAND_PREFIXES = frozenset( } ) _ASSIGNMENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") +# Env-assignment prefixes that change command lookup or code loading, so +# `LD_PRELOAD=x ls` / `PATH=. ls` run attacker code before the read-only +# utility. LD_*/DYLD_* and any *PATH are covered by the prefix/suffix check. +_AUTO_UNSAFE_ENV_ASSIGN = frozenset( + { + "IFS", + "BASH_ENV", + "ENV", + "SHELLOPTS", + "BASHOPTS", + "GLOBIGNORE", + "PROMPT_COMMAND", + "PS4", + "PYTHONSTARTUP", + "PYTHONHOME", + "NODE_OPTIONS", + "PERL5OPT", + "PERL5LIB", + "RUBYOPT", + "RUBYLIB", + # LESSOPEN/LESSCLOSE run an input preprocessor command for less. + "LESSOPEN", + "LESSCLOSE", + } +) + + +def _env_assignment_is_unsafe(name: str) -> bool: + """True if a NAME=value prefix affects command lookup/loading.""" + return ( + name in _AUTO_UNSAFE_ENV_ASSIGN + or name.startswith(("LD_", "DYLD_")) + or name.endswith("PATH") + ) + + _FIND_EXEC_FLAGS = frozenset({"-exec", "-execdir", "-ok", "-okdir"}) @@ -273,6 +310,2169 @@ def _find_blocked_commands(command: str) -> set[str]: return blocked +# ── "Approve for me" (permission_mode="auto") safety detection ────────────── +# Auto mode pauses only calls classified here as potentially unsafe. The sandbox +# and hard blocks (blocklist, rlimits) still apply at run time; this gate only +# decides prompting, and fails closed: anything not provably read-only asks. + +# Read-only commands allowed to run without confirmation in auto mode. +_AUTO_SAFE_TERMINAL_COMMANDS = frozenset( + { + "ls", + "dir", + "pwd", + # cd absent: `cd /; cat etc/passwd` escapes the workdir for a later + # relative read the path scan cannot see, so cd always asks. + "cat", + "head", + "tail", + # less/more absent: their pager escapes (+cmd, !shell, -o, LESSOPEN) can + # run a command or write a file, so they always ask. + "grep", + "egrep", + "fgrep", + "rg", + "find", + "fd", + "wc", + "sort", + "uniq", + "cut", + "tr", + "diff", + "cmp", + "file", + "stat", + "du", + "df", + # ps absent: BSD env flags (ps auxe, ps eww) dump a parent's unscrubbed + # env and can't be flag-parsed reliably, so ps always asks. + "date", + "cal", + "whoami", + "id", + "uname", + "hostname", + "uptime", + "which", + "whereis", + "type", + "basename", + "dirname", + "realpath", + "readlink", + "md5", + "md5sum", + "shasum", + "sha1sum", + "sha256sum", + "cksum", + "tree", + "printenv", + "echo", + "printf", + "true", + "false", + "test", + "[", + "seq", + "nl", + "od", + "xxd", + "hexdump", + "strings", + "column", + "paste", + "join", + "comm", + "expand", + "unexpand", + "fold", + "fmt", + "rev", + "tac", + "locale", + "arch", + "nproc", + "sw_vers", + "jq", + } +) +# Flags that turn an otherwise read-only command into a writer or executor +# (sort -o FILE, tree -o FILE, xxd -r IN OUT, find -exec/-delete/...). +_AUTO_UNSAFE_COMMAND_FLAGS = { + # --files0-from=F makes sort read the NUL-separated list of input files + # named in F, so a crafted list reads arbitrary host files indirectly. + "sort": frozenset( + {"-o", "--output", "--compress-program", "-T", "--temporary-directory", "--files0-from"} + ), + "tree": frozenset({"-o"}), + "xxd": frozenset({"-r"}), + # -c/--check makes a checksum tool read a manifest file and then read every + # path it names, so a manifest listing /etc/passwd turns `sha256sum -c list` + # into an indirect host-file read; the digest form (sha256sum file) only reads + # the named files. + "md5sum": frozenset({"-c", "--check"}), + "sha1sum": frozenset({"-c", "--check"}), + "sha256sum": frozenset({"-c", "--check"}), + "shasum": frozenset({"-c", "--check"}), + "cksum": frozenset({"-c", "--check"}), + # GNU time -o/--output/-a/--append FILE writes timing output; time is a + # wrapper, so the flag is checked before the wrapped command like env -C. + "time": frozenset({"-o", "--output", "-a", "--append"}), + # rg runs an arbitrary program per file with --pre/--hostname-bin. + "rg": frozenset({"--pre", "--hostname-bin"}), + # env -C/--chdir escapes the workdir; -S/--split-string builds a command. + "env": frozenset({"-C", "--chdir", "-S", "--split-string"}), + # ionice -p/-P/-u change the I/O priority of an already running process / + # group / user instead of forwarding to a wrapped read-only command, so a + # bare `ionice -c 3 -p ` mutates another process. ionice stays a safe + # wrapper for `ionice -c 3 `; only the process-target flags ask. + "ionice": frozenset({"-p", "-P", "-u"}), + # printf -v NAME assigns to a shell var, so `printf -v PATH %s .; ls` runs + # ./ls from the workdir. + "printf": frozenset({"-v"}), + # wc/du/find --files0-from=F read the NUL-separated list of input paths named + # in F, so a crafted list reads arbitrary host files past the literal path / + # root checks, like sort --files0-from. find spells it -files0-from (a primary). + "wc": frozenset({"--files0-from"}), + "du": frozenset({"--files0-from"}), + "find": frozenset( + { + "-exec", + "-execdir", + "-ok", + "-okdir", + "-delete", + "-fprint", + "-fprint0", + "-fprintf", + "-fls", + "-files0-from", + } + ), + # fd -x/--exec/-X/--exec-batch run a command per result; + # --base-directory/--search-path move the search root outside the workdir. + "fd": frozenset({"-x", "--exec", "-X", "--exec-batch", "--base-directory", "--search-path"}), + # date -s/--set writes the clock; display forms (+FORMAT, -d/-u/-R/-r) read. + "date": frozenset({"-s", "--set"}), + # file -C/--compile writes a compiled .mgc magic database; ident forms read. + "file": frozenset({"-C", "--compile"}), + # hostname -F/--file, -b/--boot set the hostname; display flags only read. + "hostname": frozenset({"-F", "--file", "-b", "--boot"}), +} +# Commands safe only without a mutating positional: `hostname NAME` sets the +# hostname, `date MMDDhhmm...` sets the clock (a +FORMAT token or a display +# flag's value stays read-only), so any other positional asks. +_AUTO_ARG_SENSITIVE_COMMANDS = frozenset({"hostname", "date"}) +# date display flags taking a value token (-d STRING, -r FILE, -f FILE); the +# value is not a clock-setting positional, so it is skipped. +_DATE_DISPLAY_VALUE_FLAGS = frozenset({"-d", "--date", "-r", "--reference", "-f", "--file"}) +# Commands that write their 2nd positional (uniq [INPUT [OUTPUT]], xxd [infile +# [outfile]]): the 1st file reads to stdout, but a second file positional +# overwrites it, like `sort -o`. +_AUTO_SECOND_POSITIONAL_WRITES = frozenset({"uniq", "xxd"}) +# Value-taking option flags for those commands whose argument is a separate token +# (uniq -f 2, xxd -c 16). The value must be consumed so a numeric option value is +# not miscounted as the output-file positional, and, conversely, a file that is +# literally named with digits (uniq 123 out) is still counted. +_SECOND_POSITIONAL_VALUE_FLAGS = { + "uniq": frozenset({"-f", "--skip-fields", "-s", "--skip-chars", "-w", "--check-chars"}), + "xxd": frozenset( + {"-c", "--cols", "-s", "--seek", "-l", "--len", "-g", "--groupsize", "-o", "--offset"} + ), +} +# find/fd group with (...) which resets command context, so scan every token for +# these once find/fd appears anywhere. +_AUTO_UNSAFE_FIND_LIKE_FLAGS = _AUTO_UNSAFE_COMMAND_FLAGS["find"] | _AUTO_UNSAFE_COMMAND_FLAGS["fd"] +# Recursive readers with an absolute-path target escape the workdir onto host +# files (grep -R TOKEN /home, rg TOKEN /), so they ask. +_AUTO_RECURSIVE_SEARCH = frozenset({"grep", "egrep", "fgrep", "rg", "ug", "find", "fd"}) +# Directory walkers that always recurse (tree /home, du /) read the whole host +# subtree under an absolute/tilde root, like a recursive search. ls only recurses +# with -R/--recursive, so it is gated separately when that flag is present. +_AUTO_RECURSIVE_LISTERS = frozenset({"tree", "du"}) +# Benign wrappers: safe AND forward command position to their target (checked in +# turn). sudo/su/chroot/etc. are absent, so they classify as unsafe. xargs is +# absent too: it appends arguments read from stdin that this scan never sees, so +# `echo -o out /etc/passwd | xargs sort` forwards to `sort -o out /etc/passwd` +# (a write + sensitive read) while only the allow-listed literals are visible. +_AUTO_SAFE_WRAPPERS = frozenset( + {"env", "command", "time", "timeout", "nice", "ionice", "stdbuf", "nohup"} +) + +# MCP tools whose names look read-only auto-run; anything else asks. +_AUTO_SAFE_MCP_TOOL_RE = re.compile( + r"^(get|list|search|read|fetch|query|find|describe|show|view|lookup|" + r"retrieve|count|status|info|help|check)(?:[_\-].*)?$", + re.IGNORECASE, +) +# A mutating verb anywhere in the name overrides a read-only prefix, so a +# compound name like get_or_create_issue or read_and_delete_file still asks. +_AUTO_UNSAFE_MCP_VERB_RE = re.compile( + r"(?:^|[_\-])(?:create|update|delete|remove|write|set|add|send|post|put|" + r"patch|insert|drop|kill|exec|execute|run|deploy|publish|move|rename|edit|" + r"modify|upload|replace|revoke|grant|approve|merge|close|cancel|pay|" + r"transfer|buy|sell|reset|clear|purge|destroy|terminate|revert|rollback|" + r"trigger|enable|disable|install|uninstall|restart|stop|start|" + r"save|archive|submit|commit|push|sync|register|" + r"clone|checkout|comment|fork|tag|invite|share|append|prepend|" + r"copy|duplicate|import|export|download|backup|restore|snapshot|mirror|" + r"upsert|assign|mark|subscribe|unsubscribe|reply|notify)(?:[_\-]|$)", + re.IGNORECASE, +) +# A read-named MCP tool that returns a secret is still a sensitive read, so a +# credential noun anywhere in the name (read_secret, list_tokens, +# get_credentials, fetch_api_key) asks even without a mutating verb or a path/SQL +# argument. Scoped nouns (api/access/private/... _key) avoid flagging benign +# keys like a primary_key or keyboard lookup. +_AUTO_SENSITIVE_MCP_NOUN_RE = re.compile( + r"(?:^|[_\-])(?:" + r"secret|token|credential|password|passwd|passphrase|apikey|" + r"(?:api|access|private|secret|signing|encryption|auth|session)[_\-]?keys?" + r")s?(?:[_\-]|$)", + re.IGNORECASE, +) + +# Python: modules whose import alone signals side effects auto mode should ask +# about (process spawning, network, bulk file ops, low-level memory). +_AUTO_UNSAFE_PY_MODULES = frozenset( + { + "subprocess", + "shutil", + "socket", + "ctypes", + "multiprocessing", + "pty", + "fcntl", + "requests", + "urllib", + "urllib3", + "http", + "httpx", + "aiohttp", + # huggingface_hub.hf_hub_download / snapshot_download fetch remote repo + # files over the network and write them to an on-disk cache. + "huggingface_hub", + # websockets opens a network connection; socketserver binds a listener. + "websockets", + "socketserver", + "ftplib", + "smtplib", + "telnetlib", + "paramiko", + # mail/news/rpc/browser stdlib clients open outbound connections + # (imaplib, poplib, xmlrpc.client, webbrowser.open). + "imaplib", + "poplib", + "nntplib", + "xmlrpc", + "webbrowser", + "tempfile", + # deserialization that can execute arbitrary code on load. + "pickle", + "marshal", + "shelve", + "dill", + # dbm.open(file, "c"/"n") creates files; treat the family as writers. + "dbm", + # sqlite3.connect(path) creates/mutates a database file (and runs DDL/DML + # without an open()/writer attribute), like dbm. + "sqlite3", + # runpy runs a script/module as code. + "runpy", + # ensurepip.bootstrap installs pip and venv.create builds an environment; + # both write to disk and can fetch/install packages. + "ensurepip", + "venv", + } +) +# Attribute calls that mutate the filesystem / spawn processes (os.remove, +# Path.write_text, sock.connect, ...) regardless of how the module was bound. +_AUTO_UNSAFE_PY_ATTRS = frozenset( + { + "remove", + "unlink", + "rmdir", + "removedirs", + "rename", + "renames", + "replace", + "rmtree", + "move", + "copy", + "copy2", + "copyfile", + "copytree", + "chmod", + "chown", + "system", + "popen", + "execv", + "execve", + "execl", + "execlp", + "execvp", + "spawnl", + "spawnv", + # os.startfile launches a program via its Windows association. + "startfile", + "fork", + "kill", + "killpg", + "symlink", + "link", + "mkdir", + "makedirs", + "truncate", + "touch", + "write_text", + "write_bytes", + "urlopen", + "urlretrieve", + "connect", + "bind", + "sendall", + # pathlib link creators, os node/metadata mutators, dynamic import. + "symlink_to", + "hardlink_to", + "link_to", + "mkfifo", + "mknod", + "utime", + # os.setxattr / os.removexattr mutate extended attributes, like chmod. + "setxattr", + "removexattr", + "import_module", + # loader.exec_module runs a module's code like import_module; archive + # extractall/extract write arbitrary files (zip-slip): extract takes a + # single member but an attacker-controlled member path still escapes. + "exec_module", + "extractall", + "extract", + "FileIO", + # asyncio subprocess spawners run a program past the terminal blocklist. + "create_subprocess_exec", + "create_subprocess_shell", + "subprocess_exec", + "subprocess_shell", + # asyncio outbound connections / listeners (open_connection, + # create_connection/server and unix variants), like socket.connect. + "open_connection", + "create_connection", + "create_server", + "create_unix_connection", + "create_unix_server", + # more asyncio listen/connect + UDP/raw socket helpers. + "start_server", + "start_unix_server", + "open_unix_connection", + "create_datagram_endpoint", + "sock_connect", + # os.chdir escapes the workdir; runpy helpers run arbitrary code. + "chdir", + "fchdir", + "run_path", + "run_module", + # types.FunctionType wraps a compiled code object into a callable, a + # dynamic-execution vector; pandas read_pickle deserializes (runs code). + "FunctionType", + "read_pickle", + } +) +# Pickle-backed loaders that can execute code embedded in the file; gated by +# receiver module (torch.load, joblib.load) since bare `load` is too common. +_AUTO_UNSAFE_PY_LOAD_MODULES = frozenset({"torch", "joblib", "cloudpickle"}) +# Writer methods that persist to disk without going through open() (numpy.save, +# Image.save, plt.savefig, DataFrame.to_csv, json.dump). Gated as method calls +# only, so a bare attribute reference is not mistaken for a write. +_AUTO_UNSAFE_PY_WRITE_METHODS = frozenset( + { + "save", + "savefig", + "savez", + "savez_compressed", + "savetxt", + "tofile", + "dump", + "to_csv", + "to_parquet", + "to_pickle", + "to_json", + "to_feather", + "to_hdf", + "to_excel", + "to_stata", + "to_sql", + "to_xml", + # pandas text exporters that write when given a path/buffer (to_html / + # to_markdown / to_latex mirror to_csv); to_clipboard / to_gbq persist + # off-process. to_string is omitted: it is overwhelmingly display-only. + "to_html", + "to_markdown", + "to_latex", + "to_clipboard", + "to_gbq", + "imwrite", + "imsave", + "write_image", + "write_html", + # ML persistence helpers (transformers/peft/safetensors/keras) that + # export adapters or weights to disk without an open()/writer attribute. + "save_pretrained", + "save_file", + "save_model", + "save_weights", + "save_lora", + "save_checkpoint", + # logging file handlers open a log file for write on construction (even + # default mode "a" creates); matched as attribute call and bare import. + "FileHandler", + "WatchedFileHandler", + "RotatingFileHandler", + "TimedRotatingFileHandler", + # numpy.memmap(..., mode="w+") and pandas writers create/truncate a file + # on construction, like open(..., "w"). + "memmap", + "open_memmap", + "ExcelWriter", + "HDFStore", + # pydoc.writedoc(name) writes name.html to the workdir. + "writedoc", + } +) +# Archive / compressed-file constructors taking the mode as their 2nd arg like +# open: ZipFile(name, "w") / gzip.GzipFile(name, "w") write, so gated only in +# write mode (reading a .gz is fine, so the modules are not blanket-unsafe). +_ARCHIVE_CTOR_NAMES = frozenset({"ZipFile", "TarFile", "GzipFile", "BZ2File", "LZMAFile"}) +# The stdlib module each archive constructor is imported from. +_ARCHIVE_CTOR_MODULES = { + "zipfile": "ZipFile", + "tarfile": "TarFile", + "gzip": "GzipFile", + "bz2": "BZ2File", + "lzma": "LZMAFile", +} +# Modules whose top-level open() takes the mode as its 2nd arg like builtin open, +# so `from gzip import open as gopen` binds an open alias gated on write mode. +_OPEN_ALIAS_MODULES = frozenset({"gzip", "bz2", "lzma"}) +# Builtins/itertools helpers that call their first argument once per item, so a +# writer/open alias handed to one runs without a direct call(...) site +# (list(map(open, names, modes)), starmap(np.save, ...)). filter's predicate is +# also invoked, so a writer smuggled there runs too. +_HIGHER_ORDER_INVOKERS = frozenset({"map", "filter", "starmap", "reduce"}) +_PY_WRITE_MODE_RE = re.compile(r"[wax+]") +# A file-mode literal ("w", "rb", "a+"): letters/flags only, no path chars. +# Used to tell a Path.open("w") mode from a ZipFile.open("name.txt") filename. +_PY_MODE_LITERAL_RE = re.compile(r"^[rwxa][btru+]*$") + +# Reading these off the host escapes the intent of "read-only is safe": they +# hold credentials. Path traversal (../) escapes the per-session workdir. +_SENSITIVE_PATH_RE = re.compile( + r"(?:^|[/\\])\.(?:ssh|aws|azure|gnupg|docker|kube|config/gcloud|config/gh)(?:[/\\]|$)" + r"|\.(?:netrc|npmrc|pypirc|git-credentials|env)(?:$|[/\\.\s'\"])" + r"|id_rsa|id_ed25519|id_ecdsa|id_dsa" + # Hugging Face stores the login token at ~/.cache/huggingface/token and the + # legacy ~/.huggingface/token (plus the multi-token store stored_tokens); the + # rest of that cache is model data, so only the credential files match. The + # optional leading dot covers the .huggingface dotdir form. + r"|(?:^|[/\\])\.?huggingface[/\\](?:token|stored_tokens)(?:$|[/\\.\s'\"])" + # /etc/ssh holds the host private keys (ssh_host_*_key); the whole dir is + # sensitive, not just passwd/shadow/sudoers. + r"|credentials|/etc/(?:passwd|shadow|sudoers|ssh(?:[/\\]|$))" + # Bash opens /dev/tcp/host/port and /dev/udp/host/port as network sockets, + # so a redirection to one reaches the network without the confirm prompt. + r"|/dev/(?:tcp|udp)/" + # Docker/Kubernetes secret mounts hold injected credentials. + r"|/(?:var/)?run/secrets(?:[/\\]|$)" + # procfs leaks a (possibly parent) process env/args/memory to a read, + # including the per-thread aliases under /proc//task//. The fd/ + # dir holds symlinks to a process's open files (a held credential/db file). + r"|/proc/[^/\s'\"]+/(?:task/[^/\s'\"]+/)?(?:environ|cmdline|mem|maps|fd)\b" + # A .pem/.key file (basename before the extension), not a bare ".key" + # (e.g. a jq '.key' filter). + r"|\w[\w.-]*\.(?:pem|key)(?:$|[\s'\"])", + re.IGNORECASE, +) +# A shell redirection with no following space (cat <../../notes) keeps `..` +# adjacent to `<`/`>`, so those count as leading delimiters here too. +_PARENT_TRAVERSAL_RE = re.compile(r"(?:^|[\s/\\'\"=:<>])\.\.(?:[/\\]|$|[\s'\"])") +# A sensitive directory: a dynamic segment under it (open(f"/etc/{name}")) is +# not provably safe, so fail closed when a folded path has a dynamic piece here. +_SENSITIVE_DIR_RE = re.compile( + r"/etc/|/(?:var/)?run/secrets[/\\]|(?:^|[/\\])\.(?:ssh|aws|azure|gnupg|docker|kube)[/\\]" + r"|(?:^|[/\\])\.config/(?:gcloud|gh)[/\\]", + re.IGNORECASE, +) +# Collapse /./ and repeated slashes so /etc/./passwd and /etc//passwd, which +# the OS resolves to /etc/passwd, still match the sensitive-path regex. +_REDUNDANT_SLASH_RE = re.compile(r"/\.?(?=/)") +# $name, ${name}, and operator/substring forms (${name:-x}, ${name:0:6}) all +# reference `name`; substituting the assigned value catches paths hidden behind +# a substring expansion (p=passwd; cat /etc/${p:0:6}). +_SHELL_VAR_RE = re.compile(r"\$\{(\w+)(?::[^{}]*)?\}|\$(\w+)") +# Pattern replacement (${p/X/w}, global ${p//X/w}) transforms the value before +# the path is used; apply it so p=passXd; cat /etc/${p/X/w} is scanned. +_SHELL_PARAM_REPL_RE = re.compile(r"\$\{(\w+)/(/)?([^/{}]*)/([^{}]*)\}") +# Case modification (${p^^} upper, ${p,,} lower, ${p^}/${p,} first char) also +# transforms the value, so p=PASSWD; cat /etc/${p,,} builds /etc/passwd. +_SHELL_PARAM_CASE_RE = re.compile(r"\$\{(\w+)(\^\^|,,|\^|,)\}") +# Indirect expansion ${!p} yields the value of the variable *named* by $p, so +# x=passwd; p=x; cat /etc/${!p} builds /etc/passwd. +_SHELL_PARAM_INDIRECT_RE = re.compile(r"\$\{!(\w+)\}") +_SHELL_ASSIGN_RE = re.compile(r"(?:^|[\s;&|(])([A-Za-z_]\w*)=([^\s;&|)]+)") +# Bash ANSI-C quoting ($'\x77' -> 'w') is expanded after this classifier, so +# decode $'...' bodies before the sensitive-path scan. +_ANSI_C_RE = re.compile(r"\$'((?:[^'\\]|\\.)*)'") +# Shell quotes only delimit; bash concatenates the pieces (cat /proc/x/enviro''n +# reads .../environ), so strip them before the sensitive-path scan. +_SHELL_QUOTE_RE = re.compile(r"['\"]") +# A glob bracket class [s] -> s, so .s[s]h de-obfuscates to .ssh for the scan. +_GLOB_BRACKET_RE = re.compile(r"\[([^!\]][^\]]*)\]") +# Bash POSIX character classes ([[:lower:]]) each match one char; Python fnmatch +# does not understand them, so normalize to `?` before the glob check. +_POSIX_CLASS_RE = re.compile(r"\[\[:\w+:\]\]") +# Canonical sensitive files a ? / * / [..] glob could expand to; fnmatch tests +# whether the pattern reaches one (cat /e??/passwd -> /etc/passwd). +_SENSITIVE_GLOB_TARGETS = ( + "/etc/passwd", + "/etc/shadow", + "/etc/sudoers", + "/root/.ssh/id_rsa", + "/root/.aws/credentials", + "/home/u/.ssh/id_rsa", + "/home/u/.ssh/id_ed25519", + "/home/u/.aws/credentials", + "/home/u/.netrc", + "/home/u/.git-credentials", +) +# Directories whose every file is a credential/secret; a glob resolving into one +# (cat /r?n/secrets/hf_token, cat /root/.s??/id_rsa) reads a secret even though +# the exact filename is never enumerated, so a globbed token here asks. +_SENSITIVE_GLOB_DIRS = ( + "/run/secrets", + "/var/run/secrets", + "/root/.ssh", + "/root/.aws", + "/root/.azure", + "/root/.gnupg", + "/root/.docker", + "/root/.kube", + "/root/.config/gcloud", + "/root/.config/gh", + "/home/u/.ssh", + "/home/u/.aws", + "/home/u/.azure", + "/home/u/.gnupg", + "/home/u/.docker", + "/home/u/.kube", + "/home/u/.config/gcloud", + "/home/u/.config/gh", +) +# Credential basenames a glob can reach even when the directory is not wholly +# sensitive (cat ~/.huggingface/tok?n -> token, cat ~/.netr? -> .netrc); the +# canonical-target list only covers a few fixed home paths, so match the globbed +# basename against these directly. +_SENSITIVE_GLOB_BASENAMES = frozenset( + { + "token", + "stored_tokens", + "credentials", + ".netrc", + "netrc", + ".pypirc", + ".npmrc", + ".git-credentials", + "id_rsa", + "id_ed25519", + "id_ecdsa", + "id_dsa", + "passwd", + "shadow", + # A project .env holds secrets; the literal path is gated elsewhere, so a + # glob that expands to it (cat .e?v) must be too. + ".env", + } +) +# A leading shell redirection (<, >, 2>, >>) hides the path from a plain glob +# scan (cat ]+") +# Bash brace expansion (cat /etc/pass{w,}d -> /etc/passwd /etc/passd, and the +# sequence form cat /etc/pass{w..w}d -> /etc/passwd) runs after this classifier; +# expand comma groups and .. sequences to scan each result. +_BRACE_COMMA_RE = re.compile(r"^\{([^{}]*,[^{}]*)\}$") +_BRACE_SEQ_RE = re.compile(r"^\{([^{}]+)\.\.([^{}]+)(?:\.\.(-?\d+))?\}$") +_BRACE_ANY_RE = re.compile(r"\{[^{}]*,[^{}]*\}|\{[^{}]+\.\.[^{}]+(?:\.\.-?\d+)?\}") +# Parameter expansion with a default/alternate operator (${x:-passwd}, +# ${x:+passwd}, ${x=passwd}) can synthesize a path after approval; the operand +# is substituted so the resulting path is scanned. +_SHELL_PARAM_OP_RE = re.compile(r"\$\{[A-Za-z_]\w*:?[-=+]([^{}]*)\}") + + +def _references_sensitive_path(text: str) -> bool: + """True if a command or string literal reads a credential path or escapes + the sandbox workdir via parent traversal.""" + norm = _REDUNDANT_SLASH_RE.sub("", text) + debracket = _GLOB_BRACKET_RE.sub(lambda m: m.group(1)[0], text) + return bool( + _PARENT_TRAVERSAL_RE.search(text) + or _SENSITIVE_PATH_RE.search(text) + or _SENSITIVE_PATH_RE.search(norm) + or _SENSITIVE_PATH_RE.search(debracket) + ) + + +def _pattern_matches_dir(pattern: str, target: str) -> bool: + """Segment-wise fnmatch so a glob segment does not cross a '/' boundary + (`/home/*` must not match `/home/u/.ssh`).""" + p = pattern.split("/") + t = target.split("/") + if len(p) != len(t): + return False + return all(fnmatch.fnmatch(tseg, pseg) for pseg, tseg in zip(p, t)) + + +def _glob_token_sensitive(token: str) -> bool: + """True if a single ? / * / [..] glob token could expand to a sensitive file + or a file under a secret/credential directory. Shared by the terminal scan + and the Python glob check (glob.glob('/e??/passwd')).""" + token = _REDIR_PREFIX_RE.sub("", _SHELL_QUOTE_RE.sub("", token)) + # A POSIX class ([[:lower:]]) matches one char, like `?`, but fnmatch treats + # it as a literal set; normalize so cat /etc/pass[[:lower:]]d resolves. + token = _POSIX_CLASS_RE.sub("?", token) + if not any(c in token for c in "?*["): + return False + if any(fnmatch.fnmatch(target, token) for target in _SENSITIVE_GLOB_TARGETS): + return True + # A glob that resolves to a credential basename is sensitive wherever it + # lives (cat ~/.huggingface/tok?n -> token, cat proj/.netr? -> .netrc); the + # fixed-target list only covers a handful of home paths. + base = token.rsplit("/", 1)[-1] + if any(c in base for c in "?*[") and any( + fnmatch.fnmatch(name, base) for name in _SENSITIVE_GLOB_BASENAMES + ): + return True + # A globbed directory that resolves into a secret/credential dir makes every + # file below it sensitive (cat /r?n/secrets/hf_token). + head = token.rsplit("/", 1)[0] if "/" in token else token + return any( + _pattern_matches_dir(token, d) or _pattern_matches_dir(head, d) + for d in _SENSITIVE_GLOB_DIRS + ) + + +def _glob_hits_sensitive(command: str) -> bool: + """True if any glob token in a command could expand to a sensitive file, so + `cat /e??/passwd` and `cat /r?n/secrets/hf_token` ask even without a literal + sensitive path.""" + return any( + _glob_token_sensitive(token) + for token in command.replace(";", " ").replace("|", " ").split() + ) + + +def _expand_shell_assignments(command: str) -> str: + """Best-effort substitution of `NAME=value ... $NAME`, so a sensitive path + split across an assignment and an argument (p=/etc; cat $p/passwd) is still + visible to the sensitive-path scan. Also applies pattern replacement + (p=passXd; cat /etc/${p/X/w}). Fail-open: only adds detections.""" + env = dict(_SHELL_ASSIGN_RE.findall(command)) + if not env: + return command + + def repl_pattern(m): + var, is_global, pat, rep = m.group(1), m.group(2), m.group(3), m.group(4) + if var not in env or not pat: + return m.group(0) + return env[var].replace(pat, rep) if is_global else env[var].replace(pat, rep, 1) + + def repl_case(m): + var, op = m.group(1), m.group(2) + if var not in env: + return m.group(0) + v = env[var] + if op == ",,": + return v.lower() + if op == "^^": + return v.upper() + if op == ",": + return v[:1].lower() + v[1:] + return v[:1].upper() + v[1:] + + def repl_indirect(m): + # ${!p} -> value of the variable named by $p (env[env[p]]). + pointed = env.get(m.group(1)) + return env.get(pointed, m.group(0)) if pointed is not None else m.group(0) + + command = _SHELL_PARAM_INDIRECT_RE.sub(repl_indirect, command) + command = _SHELL_PARAM_REPL_RE.sub(repl_pattern, command) + command = _SHELL_PARAM_CASE_RE.sub(repl_case, command) + return _SHELL_VAR_RE.sub(lambda m: env.get(m.group(1) or m.group(2), m.group(0)), command) + + +def _expand_param_defaults(command: str) -> str: + """Substitute the operand of a default/alternate parameter expansion + (cat /etc/pass${x:-wd} -> cat /etc/passwd), which bash applies after this + classifier. Fail-open: only adds detections.""" + return _SHELL_PARAM_OP_RE.sub(lambda m: m.group(1), command) + + +def _decode_ansi_c(command: str) -> str: + """Decode bash ANSI-C quoted words (cat $'/etc/pass\\x77d' -> cat /etc/passwd) + so an escape-obfuscated path is visible to the scan. Fail-open: only adds + detections.""" + + def dec(m): + try: + return bytes(m.group(1), "utf-8").decode("unicode_escape") + except (UnicodeDecodeError, ValueError): + return m.group(0) + + return _ANSI_C_RE.sub(dec, command) + + +def _brace_range(lo: str, hi: str, step: "str | None") -> "list[str]": + """Expand a bash sequence brace endpoint pair ({1..3}, {a..c}, {w..w}).""" + try: + istep = abs(int(step)) if step else 1 + istep = istep or 1 + if re.fullmatch(r"-?\d+", lo) and re.fullmatch(r"-?\d+", hi): + a, b = int(lo), int(hi) + rng = range(a, b + 1, istep) if a <= b else range(a, b - 1, -istep) + return [str(x) for x in rng][:64] + if len(lo) == 1 and len(hi) == 1 and lo.isalpha() and hi.isalpha(): + a, b = ord(lo), ord(hi) + rng = range(a, b + 1, istep) if a <= b else range(a, b - 1, -istep) + return [chr(x) for x in rng][:64] + except (ValueError, TypeError): + pass + return [] + + +def _brace_options(text: str) -> "list[str]": + """Options a single brace group expands to (comma list or .. sequence).""" + m = _BRACE_COMMA_RE.match(text) + if m: + return m.group(1).split(",") + m = _BRACE_SEQ_RE.match(text) + if m: + return _brace_range(m.group(1), m.group(2), m.group(3)) or [text] + return [text] + + +def _expand_braces(command: str) -> str: + """Best-effort bash brace expansion (cat /etc/pass{w,}d -> cat /etc/passwd + /etc/passd, cat /etc/pass{w..w}d -> cat /etc/passwd) so a sensitive path + split across a brace group is scanned. Bounded. Fail-open: only detects.""" + results = [command] + for _ in range(6): + if not any(_BRACE_ANY_RE.search(s) for s in results): + break + expanded = [] + for s in results: + m = _BRACE_ANY_RE.search(s) + if not m: + expanded.append(s) + continue + for opt in _brace_options(m.group(0)): + expanded.append(s[: m.start()] + opt + s[m.end() :]) + results = expanded[:64] + return " ".join(results) + + +def _mode_arg_writes(mode_node) -> bool: + """True if an AST node used as a file mode requests write/append.""" + if mode_node is None: + return False # default "r" + if isinstance(mode_node, ast.Constant) and isinstance(mode_node.value, str): + return bool(_PY_WRITE_MODE_RE.search(mode_node.value)) + return True # dynamic mode: cannot prove read-only + + +def _has_kwarg_splat(node) -> bool: + """True if the call has a ``**kwargs`` splat, which can hide a write mode.""" + return any(kw.arg is None for kw in node.keywords or []) + + +def _builtin_open_writes(node) -> bool: + """Write check for builtin ``open(file, mode)`` (mode is the 2nd arg).""" + if _has_kwarg_splat(node): + return True # **{"mode": "w"} could request a write + if any(isinstance(a, ast.Starred) for a in node.args): + return True # *("f", "w") could splat a write mode into the positionals + mode = node.args[1] if len(node.args) >= 2 else None + for kw in node.keywords or []: + if kw.arg == "mode": + mode = kw.value + return _mode_arg_writes(mode) + + +def _attr_open_writes(node) -> bool: + """Write check for ``x.open(...)`` (e.g. ``Path.open(mode)`` where mode is + the 1st arg). Only a mode-looking string is read as the mode, so a + ``ZipFile.open("name.txt")`` read is not mistaken for a write.""" + if _has_kwarg_splat(node): + return True # **{"mode": "w"} could request a write + for kw in node.keywords or []: + if kw.arg == "mode": + return _mode_arg_writes(kw.value) + if node.args: + first = node.args[0] + if isinstance(first, ast.Constant) and isinstance(first.value, str): + if _PY_MODE_LITERAL_RE.match(first.value): + return bool(_PY_WRITE_MODE_RE.search(first.value)) + # A 2nd positional arg is either a mode (x.open(name, "w")) or + # os.open(path, O_CREAT) flags via an alias: honor a string mode, + # otherwise cannot prove read-only, so ask. + if len(node.args) >= 2: + second = node.args[1] + if isinstance(second, ast.Constant) and isinstance(second.value, str): + return _mode_arg_writes(second) + return True + return False + return True # dynamic first arg: cannot prove read-only + return False # no args: read + + +_PATH_CTORS = ( + "Path", + "PurePath", + "PurePosixPath", + "PureWindowsPath", + "PosixPath", + "WindowsPath", +) +# Deterministic path pass-through/normalizer calls that return the same location +# (os.path.abspath('/etc') -> /etc, Path('/etc').resolve() -> /etc), so folding +# through them keeps a sensitive root visible to the scan. +_PATH_PASSTHROUGH_ATTRS = frozenset( + {"abspath", "normpath", "realpath", "expanduser", "expandvars", "resolve", "absolute"} +) +# pathlib methods that rewrite only the final path component, so the sensitive +# target is never spelled out as a literal (Path('/etc/x').with_name('passwd') +# -> /etc/passwd). Folded below so the rewritten path is still scanned. +_PATH_NAME_REWRITES = frozenset({"with_name", "with_stem", "with_suffix"}) +# Mapping-style %-format conversion specifier: %(name)s / %(n)5.2f. Used to fold +# '/etc/%(f)s' % {'f': 'passwd'} to /etc/passwd (a dynamic value becomes NUL). +_PERCENT_NAMED_RE = re.compile(r"%\((\w+)\)[-#0 +]*\d*(?:\.\d+)?[a-zA-Z]") + + +def _folded_path( + node, + literals = None, + ctors = None, + join_names = None, +) -> "str | None": + """Best-effort value of a path built from string literals, so a sensitive + path assembled from pieces (os.path.join('/etc', 'passwd'), '/etc'+'/passwd', + Path('/etc') / 'passwd', f'/proc/{pid}/environ', f'/etc/{name}') is still + visible to the scan. A dynamic piece becomes NUL, a non-slash placeholder, + so a dynamic segment under a sensitive dir (/etc/NUL) is still detectable. + ``literals`` maps names bound to string literals (base = '/etc'); ``ctors`` + is the set of pathlib constructor names (incl. import aliases); ``join_names`` + are bare names bound to os.path.join (from os.path import join).""" + literals = literals or {} + ctors = ctors or _PATH_CTORS + join_names = join_names or frozenset() + + def fold(node) -> "str | None": + if isinstance(node, ast.Constant) and isinstance(node.value, (str, bytes)): + # bytes paths are valid too (open(b'/etc/passwd')); decode for scan. + return ( + node.value.decode("latin-1", "ignore") + if isinstance(node.value, bytes) + else node.value + ) + if isinstance(node, ast.Name): + return literals.get(node.id) + if isinstance(node, ast.Attribute) and node.attr in ("parent", "parents"): + # A pathlib .parent/.parents walks above the current dir, escaping + # the per-session workdir without a literal '..'; mark it so a read + # folds to unsafe (\x02 is a non-slash escape sentinel). + return "\x02" + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Attribute) + and (node.value.attr == "parents") + ): + return "\x02" # Path(...).parents[1] + if isinstance(node, ast.JoinedStr): + return "".join( + v.value + if isinstance(v, ast.Constant) and isinstance(v.value, str) + else (fold(v.value) or "\x00") + if isinstance(v, ast.FormattedValue) + else "\x00" + for v in node.values + ) + if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Div)): + left = fold(node.left) + right = fold(node.right) + left = "\x00" if left is None else left + right = "\x00" if right is None else right + # Path('/etc') / 'passwd' joins with a separator; '+' concatenates. + return left + "/" + right if isinstance(node.op, ast.Div) else left + right + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Mod): + # Old-style formatting: '%s/%s' % ('/etc', 'passwd') -> /etc/passwd. + template = fold(node.left) + if template is not None and "%" in template: + rhs = node.right + if "%(" in template: + # Mapping-style: '/etc/%(f)s' % {'f': 'passwd'} -> /etc/passwd. + # A literal dict resolves each name; an unresolved value or a + # non-literal mapping leaves the NUL marker so /etc/ + # still fails closed under a sensitive dir. + mapping: "dict[str, str]" = {} + if isinstance(rhs, ast.Dict): + for k, v in zip(rhs.keys, rhs.values): + if isinstance(k, ast.Constant) and isinstance(k.value, str): + fv = fold(v) + mapping[k.value] = fv if fv is not None else "\x00" + return _PERCENT_NAMED_RE.sub( + lambda m: mapping.get(m.group(1), "\x00"), template + ) + if isinstance(rhs, ast.Tuple): + args = tuple((fold(e) or "\x00") for e in rhs.elts) + else: + single = fold(rhs) + args = (single if single is not None else "\x00",) + try: + return template % args + except (TypeError, ValueError, KeyError): + return None + return None + if isinstance(node, ast.Call): + func = node.func + if isinstance(func, ast.Attribute) and func.attr == "joinpath": + # Path('/etc').joinpath('passwd') -> receiver and args are pieces. + base = fold(func.value) + parts = [base if base is not None else "\x00"] + parts += [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + if isinstance(func, ast.Attribute) and func.attr in ("glob", "rglob", "iglob"): + # Path('/etc').glob('passw?') -> the receiver dir joined with the + # glob pattern; _glob_token_sensitive then tests /etc/passw?. + base = fold(func.value) + pattern = fold(node.args[0]) if node.args else "\x00" + return (base if base is not None else "\x00") + "/" + (pattern or "\x00") + if isinstance(func, ast.Attribute) and func.attr in _PATH_NAME_REWRITES: + # Path('/etc/x').with_name('passwd') -> /etc/passwd; with_stem / + # with_suffix rewrite only the final component. Fold to the + # rewritten path so a sensitive target that no literal spells out + # is still caught. An unresolved receiver stays None (untracked, + # like a bare variable), and a dynamic arg becomes the NUL marker. + base = fold(func.value) + if base is None: + return None + arg = fold(node.args[0]) if node.args else None + arg = "\x00" if arg is None else arg + idx = base.rfind("/") + head = base[: idx + 1] if idx >= 0 else "" + name = base[idx + 1 :] if idx >= 0 else base + dot = name.rfind(".") + stem = name[:dot] if dot > 0 else name + suffix = name[dot:] if dot > 0 else "" + if func.attr == "with_name": + name = arg + elif func.attr == "with_stem": + name = arg + suffix + else: # with_suffix + name = stem + arg + return head + name + if isinstance(func, ast.Attribute) and func.attr in _PATH_PASSTHROUGH_ATTRS: + # Deterministic normalizers keep the same path: os.path.abspath( + # '/etc') -> /etc, Path('/etc').resolve() -> /etc. When called with + # a path arg fold it, else fold the receiver (Path method form). + return fold(node.args[0]) if node.args else fold(func.value) + if isinstance(func, ast.Attribute) and func.attr == "join": + # str.join has the separator as the receiver and the pieces in + # one iterable arg ("".join(['/etc', '/passwd']) -> /etc/passwd); + # tell it apart from os.path.join(*pieces). + sep = fold(func.value) + if ( + sep is not None + and len(node.args) == 1 + and isinstance(node.args[0], (ast.List, ast.Tuple)) + ): + pieces = [(fold(e) or "\x00") for e in node.args[0].elts] + return sep.join(pieces) + parts = [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + # A bare os.path.join alias (from os.path import join): join(*pieces). + if isinstance(func, ast.Name) and func.id in join_names: + parts = [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + # A bare/qualified/aliased pathlib constructor (Path(...), P(...)). + if (isinstance(func, ast.Attribute) and func.attr in ctors) or ( + isinstance(func, ast.Name) and func.id in ctors + ): + parts = [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + # '/etc/{}'.format('passwd') -> /etc/passwd (literal template + args). + if isinstance(func, ast.Attribute) and func.attr == "format": + template = fold(func.value) + if template is not None and "{" in template: + parts = [] + for a in node.args: + if isinstance(a, ast.Constant): + parts.append(str(a.value)) + else: + folded = fold(a) + parts.append("\x00" if folded is None else folded) + try: + return template.format(*parts) + except (IndexError, KeyError, ValueError): + return None + return None + + return fold(node) + + +def _dynamic_name_hits_sensitive(folded) -> bool: + """True if a folded path with a dynamic piece (NUL) inside a path segment + could spell a credential target, e.g. open('/et' + chr(99) + '/passwd') + folds to '/et\\x00/passwd'. NUL matches any run of non-separator chars so the + dynamic split of a sensitive name resolves, while an all-dynamic ('\\x00\\x00') + or segment-spanning ('\\x00/\\x00') path cannot form a single credential name + and stays safe.""" + if not folded or "\x00" not in folded: + return False + pattern = "".join(r"[^/\\]*" if ch == "\x00" else re.escape(ch) for ch in folded) + try: + rx = re.compile(pattern + r"\Z") + except re.error: + return True # pathological pattern: fail closed + return any(rx.match(t) for t in _SENSITIVE_GLOB_TARGETS) + + +def _folded_is_sensitive(folded) -> bool: + """A folded path is sensitive if it names a credential file, has a dynamic + segment (NUL) directly under a sensitive directory (/etc/NUL), walks out of + the sandbox via a pathlib .parent/.parents escape (\\x02), or is a glob that + could resolve to a credential path (glob.glob('/e??/passwd')).""" + if not folded: + return False + return ( + "\x02" in folded + or _references_sensitive_path(folded) + or ("\x00" in folded and bool(_SENSITIVE_DIR_RE.search(folded))) + # A dynamic segment (NUL) can be the "/" forming a sensitive root: + # open(os.sep + "etc/passwd") folds to "\x00etc/passwd", so re-scan with + # NUL as "/" (a benign "\x00data/file" -> "/data/file" stays safe). + or ("\x00" in folded and _references_sensitive_path(folded.replace("\x00", "/"))) + # A dynamic piece can also sit INSIDE a sensitive name: open('/et' + + # chr(99) + '/passwd') folds to "/et\x00/passwd", which none of the above + # catch. Match the literals around each NUL against a credential target, + # treating NUL as "any run of non-separator chars" so /et/passwd + # resolves while an all-dynamic ("\x00\x00" from 1 + 1) or segment-spanning + # ("\x00/\x00" from a + '/' + b) path stays safe. + or _dynamic_name_hits_sensitive(folded) + or _glob_token_sensitive(folded) + ) + + +def _terminal_is_potentially_unsafe(command: str) -> bool: + """Classify a terminal command for auto mode (fail closed).""" + if not command or not command.strip(): + return False + # Redirections and substitutions can hide writes or nested commands; a + # quoted ">" false-positives into a prompt, which is the safe direction. + if ">" in command or "`" in command or "$(" in command or "<(" in command: + return True + # Reads that escape the sandbox workdir (../) or hit credential paths are + # not "safe" reads; ask before running them. Strip shell quotes/backslash + # escapes and expand NAME=value prefixes first so `cat /proc/$PPID/enviro''n`, + # `cat /et\c/passwd`, and `p="/proc/$PPID"; cat $p/environ` are caught too. + stripped = _SHELL_QUOTE_RE.sub("", command).replace("\\", "") + # Bash applies brace/parameter/ANSI-C expansion after this classifier, so a + # path split across a brace group (/etc/pass{w,}d), a default/substring param + # (${x:-wd}, ${p:0:6}), or an escape ($'...') is invisible to the raw scan; + # expand first (ANSI-C decoded from the raw command, before backslash strip). + candidates = [] + for c in (command, stripped, _decode_ansi_c(command)): + c_param = _expand_param_defaults(c) + candidates.extend((c, c_param, _expand_braces(c_param), _expand_shell_assignments(c_param))) + # Run both the literal and glob-sensitive scans over every candidate, so a + # brace-expanded glob (cat /e{t,}c/pass?d -> /etc/pass?d) is caught. + if any(_glob_hits_sensitive(c) or _references_sensitive_path(c) for c in candidates): + return True + # Newlines (and CR) separate commands in a shell but read as plain + # whitespace to shlex, which would demote "ls\nrm x" to argument position. + command = command.replace("\r\n", ";").replace("\n", ";").replace("\r", ";") + try: + lexer = shlex.shlex(command, posix = True, punctuation_chars = ";&|()") + lexer.whitespace_split = True + tokens = list(lexer) + except ValueError: + return True + # A root can also hide behind an assignment (p=/; grep -R TOKEN $p) or a + # default parameter (grep -R TOKEN ${root:-/home}); re-lex the fully expanded + # command so the find/fd and recursive-search scans see the resolved token. + expanded_command = _expand_shell_assignments(_expand_param_defaults(command)) + if expanded_command != command: + try: + elexer = shlex.shlex(expanded_command, posix = True, punctuation_chars = ";&|()") + elexer.whitespace_split = True + scan_tokens = list(elexer) + except ValueError: + return True + else: + scan_tokens = tokens + # find/fd group with (...) which resets command context, so a trailing + # -delete/-exec could slip past; scan every token when find/fd appears. + if any(os.path.basename(t.strip(";&|()`{}")).lower() in ("find", "fd") for t in scan_tokens): + if any(t.split("=", 1)[0] in _AUTO_UNSAFE_FIND_LIKE_FLAGS for t in scan_tokens): + return True + # A recursive reader rooted outside the sandbox reads host files (grep -R + # TOKEN /home, rg TOKEN /, grep -R TOKEN ~root, p=/; grep -R TOKEN $p, and + # the always-recursive walkers tree /home / du /); ask. Bash expands + # ~/~user to a home dir after this decision, so a tilde root is a sandbox + # escape too. A path-qualified command token starts with "/" as well, but + # that already asks below. + if any(t.startswith("/") or t.startswith("~") for t in scan_tokens): + token_bases = [os.path.basename(t.strip(";&|()`{}")).lower() for t in tokens] + if any(b in _AUTO_RECURSIVE_SEARCH or b in _AUTO_RECURSIVE_LISTERS for b in token_bases): + return True + # ls only walks the whole subtree with -R/--recursive (ls -R /home, + # ls -laR /); a non-recursive ls /home lists one level and stays here. + if "ls" in token_bases and any( + t.split("=", 1)[0] in ("-R", "--recursive") + or (t[:1] == "-" and t[:2] != "--" and "=" not in t and "R" in t[1:]) + for t in tokens + ): + return True + expect_command = True + prefix_pending = False + current_command = "" + positional_args = 0 + pending_flag_value = False + for token in tokens: + # Runs of punctuation (";;", ";&") lex as one token; any token made + # purely of separator characters still separates commands. + if ( + token in _SHELL_SEPARATORS + or token in _SHELL_KEYWORDS_AS_SEP + or not set(token) - set(";&|()") + ): + expect_command = True + prefix_pending = False + current_command = "" + positional_args = 0 + pending_flag_value = False + continue + if token.startswith("-"): + # A write/exec flag on an otherwise read-only command asks + # (sort -o, tree -o, xxd -r, find -exec/-delete/...). Match + # "--output=x", an attached short option "-o/tmp/out", and a short + # option bundled in a cluster (sort -uo out => -u -o). + flag_head = token.split("=", 1)[0] + cluster = token[1:] if token[:2] != "--" and "=" not in token else "" + # GNU tools accept unambiguous abbreviations of a long option, so + # `sort --out=` reaches --output and `env --ch=/` reaches --chdir; + # a "--x" prefix of an unsafe long flag fails closed. + is_long_abbrev = flag_head.startswith("--") and len(flag_head) > 2 + for uf in _AUTO_UNSAFE_COMMAND_FLAGS.get(current_command, ()): + if flag_head == uf or (len(uf) == 2 and (token.startswith(uf) or uf[1] in cluster)): + return True + if is_long_abbrev and uf.startswith("--") and uf.startswith(flag_head): + return True + # A flag that takes a following value (date -d STRING / -r FILE; + # uniq -f N; xxd -c N) so the value token is not mistaken for a + # clock-setting positional or an output-file positional. + pending_flag_value = "=" not in token and ( + (current_command == "date" and flag_head in _DATE_DISPLAY_VALUE_FLAGS) + or flag_head in _SECOND_POSITIONAL_VALUE_FLAGS.get(current_command, ()) + ) + if not prefix_pending: + expect_command = False + continue + if not expect_command: + raw_pos = token.strip(";&|()`{}") + # uniq [INPUT [OUTPUT]] writes its second file positional; count file + # positionals and ask on the second one. A preceding option's value + # (uniq -f 2) is consumed via pending_flag_value, so a file literally + # named with digits (uniq 123 out) is still counted. + if current_command in _AUTO_SECOND_POSITIONAL_WRITES: + if pending_flag_value: + pending_flag_value = False + elif raw_pos: + positional_args += 1 + if positional_args >= 2: + return True + # hostname NAME sets the hostname; date sets the clock. A + # positional past a display flag's value therefore mutates state and + # asks (date's +FORMAT display token stays read-only). + elif current_command in _AUTO_ARG_SENSITIVE_COMMANDS: + if pending_flag_value: + pending_flag_value = False + elif raw_pos and not (current_command == "date" and raw_pos.startswith("+")): + return True + continue + if _ASSIGNMENT_RE.match(token): + # Benign NAME=value prefixes are skipped, but ones that change + # command lookup/loading (PATH, LD_PRELOAD, ...) fail closed. + if _env_assignment_is_unsafe(token.split("=", 1)[0]): + return True + continue + if prefix_pending and token.lstrip("-").isdigit(): + continue + raw = token.strip(";&|()`{}") + # A path-qualified command (./ls, /tmp/cat) is an arbitrary executable, + # not the trusted system utility its basename matches; ask first. + if "/" in raw or "\\" in raw: + return True + base = os.path.basename(raw).lower() + stem, ext = os.path.splitext(base) + if ext in {".exe", ".com", ".bat", ".cmd"}: + base = stem + if base in _AUTO_SAFE_WRAPPERS: + prefix_pending = True + # Track the wrapper so its own flags (env --chdir) are checked; + # the real command overwrites this when it is reached. + current_command = base + pending_flag_value = False + continue + if base not in _AUTO_SAFE_TERMINAL_COMMANDS: + return True + current_command = base + expect_command = False + prefix_pending = False + positional_args = 0 + pending_flag_value = False + return False + + +def _python_is_potentially_unsafe(code: str) -> bool: + """Classify python-tool code for auto mode (fail closed).""" + if not code or not code.strip(): + return False + # Anything the sandbox's static analysis already objects to would be + # refused at execution time; surface it as a confirmation first. + if _check_code_safety(code) is not None: + return True + try: + tree = ast.parse(code) + except SyntaxError: + return False # runs into a normal traceback; nothing to guard + # Names bound to the builtin open (f = open; from builtins import open as f; + # f, _ = (open, print)) so an aliased writer call is still checked below. + # builtins_aliases tracks `import builtins [as b]` for builtins.exec/eval. + open_aliases = {"open"} + # Attribute names bound to open (box.f = open), so a later box.f('out', 'w') + # write is still gated even though the callable is an attribute, not a name. + attr_open_aliases: "set[str]" = set() + builtins_aliases = {"builtins", "__builtins__"} + # Names bound to a dynamic lookup (rm = getattr(os, "remove"); + # f = globals()["open"]) whose calls cannot be proven read-only, so they + # fail closed. + dynamic_aliases = set() + # Names bound to a dynamic-code builtin, including aliased ones + # (from builtins import eval as e; e = builtins.exec), so a call or + # reference through the alias fails closed too. compile() builds a code + # object that FunctionType/exec can then run. + code_exec_aliases = {"exec", "eval", "__import__", "breakpoint", "compile"} + # Names bound to a string literal (base = '/etc'), so a sensitive path + # split through a variable (base + '/passwd') folds and is caught. + literal_str_vars: "dict[str, str]" = {} + # Pathlib constructor names incl. import aliases (from pathlib import Path as + # P), os.path.join names bound directly (from os.path import join as j), and + # writer functions imported as bare names (from numpy import save). + path_ctor_aliases = set(_PATH_CTORS) + pathjoin_aliases: "set[str]" = set() + writer_aliases: "set[str]" = set() + # Module names bound to os/posix (import os as o), so o.open(...) is still + # recognized as the low-level create/write that os.open is. + os_aliases = {"os", "posix"} + # Module names bound to a pickle-backed loader (import torch as t), so + # t.load(...) is still gated as a code-executing deserialize. + load_module_aliases = set(_AUTO_UNSAFE_PY_LOAD_MODULES) + # Names bound to the builtin getattr (g = getattr), so a dynamic lookup + # aliased through it (rm = g(os, "remove"); rm("f")) still fails closed. + getattr_aliases = {"getattr"} + # Names bound to functools.partial, so a partial that wraps open/a writer + # (w = partial(open, mode="w"); w("out.txt")) fails closed when w is called. + partial_aliases: "set[str]" = set() + # Archive constructors imported bare (from zipfile import ZipFile), so + # ZipFile(name, "w") is gated like the zipfile.ZipFile attribute call. + archive_ctor_aliases: "set[str]" = set() + # operator.methodcaller("write_text") is dynamic dispatch, like getattr. + operator_aliases = {"operator"} + methodcaller_aliases: "set[str]" = set() + # logging.basicConfig(filename=...) opens a log file for write. + basicconfig_aliases: "set[str]" = set() + # fileinput.input(..., inplace=True) rewrites a file in place. + fileinput_aliases = {"fileinput"} + # Higher-order invokers (map/filter/starmap/reduce) call their first arg, so + # one handed a writer (map(open, ...)) writes without a direct open() site. + # Track aliases (m = map; from itertools import starmap as sm) so an aliased + # invoker is still checked; the write-callable gate keeps map(len, ...) safe. + invoker_aliases = set(_HIGHER_ORDER_INVOKERS) + + def _is_dynamic_namespace(node) -> bool: + # A namespace mapping whose .get/.pop/.setdefault (or subscript) can return + # open/eval/a mutator: globals()/locals()/vars(...), any X.__dict__, + # __builtins__, sys.modules. Looking a name up through one is as dynamic as + # getattr, so a value fetched from it fails closed. + if isinstance(node, ast.Attribute): + if node.attr == "__dict__": + return True + return ( + node.attr == "modules" + and isinstance(node.value, ast.Name) + and node.value.id == "sys" + ) + if isinstance(node, ast.Name): + return node.id in builtins_aliases + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + return node.func.id in ("globals", "locals", "vars") + return False + + def _methodcaller_writes(call) -> bool: + # operator.methodcaller("write_text", ...) / methodcaller(name): unsafe + # when the method name is a known writer/mutator, or non-constant (cannot + # be proven read-only). + if not call.args: + return False + first = call.args[0] + if not (isinstance(first, ast.Constant) and isinstance(first.value, str)): + return True + return first.value in _AUTO_UNSAFE_PY_ATTRS or first.value in _AUTO_UNSAFE_PY_WRITE_METHODS + + def _fileinput_inplace(call) -> bool: + # fileinput.input(..., inplace=True) opens each file for in-place rewrite. + if _has_kwarg_splat(call): + return True + for kw in call.keywords or []: + if kw.arg == "inplace": + v = kw.value + if isinstance(v, ast.Constant): + return bool(v.value) + return True # dynamic inplace flag: cannot prove read-only + return False + + def _basicconfig_writes(call) -> bool: + # logging.basicConfig(filename=...) creates/opens a log file for writing. + if _has_kwarg_splat(call): + return True + return any(kw.arg == "filename" for kw in call.keywords or []) + + def _wraps_write_callable(arg) -> bool: + # The callable a partial wraps (partial(open, ...)); True when calling it + # could create/overwrite a file or resolve a dynamic/mutating function. + if isinstance(arg, ast.Name): + return ( + arg.id in open_aliases + or arg.id in dynamic_aliases + or arg.id in code_exec_aliases + or arg.id in getattr_aliases + or arg.id in writer_aliases + or arg.id in archive_ctor_aliases + ) + if isinstance(arg, ast.Attribute): + return ( + arg.attr == "open" + or arg.attr in _AUTO_UNSAFE_PY_ATTRS + or arg.attr in _AUTO_UNSAFE_PY_WRITE_METHODS + or arg.attr in _ARCHIVE_CTOR_NAMES + ) + return False + + def _passed_write_callable(arg) -> bool: + # A concrete write callable handed as an argument to another call: a + # name bound to open / a writer / an archive constructor, or an + # attribute reference to a writer method / mutating os attr / archive + # ctor / .open. Unlike _wraps_write_callable this omits the fail-closed + # dynamic / getattr / code-exec poison aliases, which are already gated + # where they are *called* and would over-trigger when a benign alias is + # merely passed or printed (print(getattr(o, 'name'))). + if isinstance(arg, ast.Name): + return ( + arg.id in open_aliases or arg.id in writer_aliases or arg.id in archive_ctor_aliases + ) + if isinstance(arg, ast.Attribute): + return ( + arg.attr == "open" + or arg.attr in _AUTO_UNSAFE_PY_ATTRS + or arg.attr in _AUTO_UNSAFE_PY_WRITE_METHODS + or arg.attr in _ARCHIVE_CTOR_NAMES + ) + return False + + # Names bound more than once cannot be folded to a single literal: this scan + # visits every assignment before any call is checked, so a later benign + # reassignment (base = '/etc'; open(base + '/passwd'); base = 'data') would + # otherwise mask the earlier sensitive value and auto-approve. Count every + # binding target up front and poison multiply-bound names to the escape + # sentinel so any path folded from them fails closed (asks) instead. + assign_counts: "dict[str, int]" = {} + for node in ast.walk(tree): + binding_targets = [] + if isinstance(node, ast.Assign): + binding_targets = node.targets + elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): + binding_targets = [node.target] + for target in binding_targets: + for sub in ast.walk(target): + if isinstance(sub, ast.Name): + assign_counts[sub.id] = assign_counts.get(sub.id, 0) + 1 + multi_assigned_names = {name for name, count in assign_counts.items() if count > 1} + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "builtins": + builtins_aliases.add(alias.asname or "builtins") + elif alias.name in ("os", "posix"): + os_aliases.add(alias.asname or alias.name) + elif alias.name in _AUTO_UNSAFE_PY_LOAD_MODULES: + load_module_aliases.add(alias.asname or alias.name) + elif alias.name == "operator": + operator_aliases.add(alias.asname or "operator") + elif alias.name == "fileinput": + fileinput_aliases.add(alias.asname or "fileinput") + elif isinstance(node, ast.ImportFrom): + if node.module == "operator": + for alias in node.names: + if alias.name == "methodcaller": + methodcaller_aliases.add(alias.asname or "methodcaller") + if node.module == "logging": + for alias in node.names: + if alias.name == "basicConfig": + basicconfig_aliases.add(alias.asname or "basicConfig") + if node.module == "builtins": + for alias in node.names: + if alias.name == "open": + open_aliases.add(alias.asname or "open") + elif alias.name in code_exec_aliases: + code_exec_aliases.add(alias.asname or alias.name) + if node.module in _OPEN_ALIAS_MODULES: + for alias in node.names: + if alias.name == "open": + # gzip/bz2/lzma open(file, mode) writes on "w"/"a"/"x", + # mode in the 2nd arg like builtin open. + open_aliases.add(alias.asname or "open") + if node.module == "pathlib": + for alias in node.names: + if alias.name in _PATH_CTORS: + path_ctor_aliases.add(alias.asname or alias.name) + if node.module in ("os.path", "posixpath", "ntpath"): + for alias in node.names: + if alias.name == "join": + pathjoin_aliases.add(alias.asname or "join") + if node.module == "functools": + for alias in node.names: + if alias.name == "partial": + partial_aliases.add(alias.asname or "partial") + if node.module in _ARCHIVE_CTOR_MODULES: + _ctor = _ARCHIVE_CTOR_MODULES[node.module] + for alias in node.names: + if alias.name == _ctor: + archive_ctor_aliases.add(alias.asname or _ctor) + for alias in node.names: + if alias.name in _AUTO_UNSAFE_PY_WRITE_METHODS: + writer_aliases.add(alias.asname or alias.name) + # from itertools import starmap as sm / from functools import + # reduce as r: an aliased higher-order invoker. + if alias.name in _HIGHER_ORDER_INVOKERS: + invoker_aliases.add(alias.asname or alias.name) + elif isinstance(node, (ast.Assign, ast.AnnAssign)) and node.value is not None: + value = node.value + # AnnAssign (f: object = open) has a single target, no destructuring. + if isinstance(node, ast.AnnAssign): + assign_targets = [node.target] + else: + assign_targets = node.targets + targets = [t.id for t in assign_targets if isinstance(t, ast.Name)] + attr_targets = [t.attr for t in assign_targets if isinstance(t, ast.Attribute)] + if isinstance(value, ast.Name) and value.id in open_aliases: + open_aliases.update(targets) + attr_open_aliases.update(attr_targets) # box.f = open + elif isinstance(value, ast.Name) and value.id in getattr_aliases: + getattr_aliases.update(targets) # g = getattr + elif isinstance(value, ast.Name) and value.id in partial_aliases: + partial_aliases.update(targets) # p = partial + elif isinstance(value, ast.Name) and value.id in writer_aliases: + writer_aliases.update(targets) # s = save (numpy save alias) + elif isinstance(value, ast.Name) and value.id in archive_ctor_aliases: + archive_ctor_aliases.update(targets) # z = ZipFile + elif isinstance(value, ast.Name) and value.id in invoker_aliases: + invoker_aliases.update(targets) # m = map + elif isinstance(value, ast.Name) and value.id in path_ctor_aliases: + path_ctor_aliases.update(targets) # P = Path + elif isinstance(value, ast.Name) and value.id in pathjoin_aliases: + pathjoin_aliases.update(targets) # j = join + elif isinstance(value, ast.Attribute) and value.attr == "join": + pathjoin_aliases.update(targets) # j = os.path.join + elif isinstance(value, ast.Attribute) and value.attr in _PATH_CTORS: + path_ctor_aliases.update(targets) # P = pathlib.Path + elif ( + isinstance(value, ast.Attribute) + and value.attr == "open" + and isinstance(value.value, ast.Name) + and value.value.id in builtins_aliases + ): + open_aliases.update(targets) # f = builtins.open + elif ( + isinstance(value, ast.Attribute) + and value.attr in code_exec_aliases + and isinstance(value.value, ast.Name) + and value.value.id in builtins_aliases + ): + code_exec_aliases.update(targets) # e = builtins.eval + elif isinstance(value, ast.Attribute) and value.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + writer_aliases.update(targets) # s = np.save + elif isinstance(value, ast.Attribute) and value.attr == "open": + # A captured .open bound method (p = Path('out').open) opens a file + # on any call; its mode position varies (Path.open mode is 1st arg, + # builtin open's is 2nd), so fail closed on the call rather than + # guess the write mode. + dynamic_aliases.update(targets) # p = Path('out').open; p('w') + elif isinstance(value, ast.Attribute) and value.attr in _ARCHIVE_CTOR_NAMES: + archive_ctor_aliases.update(targets) # z = zipfile.ZipFile + elif isinstance(value, ast.Subscript): + dynamic_aliases.update(targets) # f = globals()["open"] + elif ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Name) + and value.func.id in getattr_aliases + ): + dynamic_aliases.update(targets) # rm = getattr(os, "remove") / g(...) + elif ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Attribute) + and value.func.attr in ("get", "pop", "setdefault") + and _is_dynamic_namespace(value.func.value) + ): + # f = __builtins__.__dict__.get("open") / globals().get("open"): + # a namespace lookup can return open/eval, so poison like getattr. + dynamic_aliases.update(targets) + elif ( + isinstance(value, ast.Call) + and ( + (isinstance(value.func, ast.Name) and value.func.id in partial_aliases) + or (isinstance(value.func, ast.Attribute) and value.func.attr == "partial") + ) + and value.args + and _wraps_write_callable(value.args[0]) + ): + dynamic_aliases.update(targets) # w = partial(open, mode="w") + elif ( + isinstance(value, ast.Call) + and ( + (isinstance(value.func, ast.Name) and value.func.id in methodcaller_aliases) + or ( + isinstance(value.func, ast.Attribute) + and value.func.attr == "methodcaller" + and isinstance(value.func.value, ast.Name) + and value.func.value.id in operator_aliases + ) + ) + and _methodcaller_writes(value) + ): + dynamic_aliases.update(targets) # w = methodcaller("write_text", ...) + elif isinstance(value, ast.Constant) and isinstance(value.value, str): + # base = '/etc' -> resolve base in a later folded path. A name + # bound more than once is poisoned (\x02) so it fails closed. + for t in targets: + literal_str_vars[t] = "\x02" if t in multi_assigned_names else value.value + elif isinstance(value, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)): + # p = Path('/etc'); q = p; r = os.path.join('/etc','x'): record a + # fully-literal folded path so a later reuse (p / 'passwd') folds. + folded = _folded_path(value, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + if folded is not None and "\x00" not in folded and "\x02" not in folded: + for t in targets: + literal_str_vars[t] = "\x02" if t in multi_assigned_names else folded + elif isinstance(value, (ast.Tuple, ast.List)): + # Destructuring binds each element like a single assignment, so an + # aliased callable (f, _ = (open, print)) AND a string / path + # literal (base, leaf = ('/etc', 'passwd')) both propagate; without + # the latter a path folded from base/leaf would miss the sensitive + # target and auto-approve. + for target in assign_targets: + if isinstance(target, (ast.Tuple, ast.List)) and len(target.elts) == len( + value.elts + ): + for tgt_el, val_el in zip(target.elts, value.elts): + if not isinstance(tgt_el, ast.Name): + continue + tid = tgt_el.id + if isinstance(val_el, ast.Name) and val_el.id in open_aliases: + open_aliases.add(tid) + elif isinstance(val_el, ast.Name) and val_el.id in getattr_aliases: + getattr_aliases.add(tid) + elif isinstance(val_el, ast.Name) and val_el.id in partial_aliases: + partial_aliases.add(tid) + elif isinstance(val_el, ast.Name) and val_el.id in writer_aliases: + writer_aliases.add(tid) # s, _ = (save, 1) + elif isinstance(val_el, ast.Name) and val_el.id in archive_ctor_aliases: + archive_ctor_aliases.add(tid) # z, _ = (ZipFile, 1) + elif isinstance(val_el, ast.Constant) and isinstance(val_el.value, str): + literal_str_vars[tid] = ( + "\x02" if tid in multi_assigned_names else val_el.value + ) + elif isinstance(val_el, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)): + folded = _folded_path( + val_el, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) + if ( + folded is not None + and "\x00" not in folded + and "\x02" not in folded + ): + literal_str_vars[tid] = ( + "\x02" if tid in multi_assigned_names else folded + ) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + # A callable captured as a parameter default (def f(o=open): o('x','w')) + # binds that parameter to the same alias set, so a later call through + # the parameter is still gated. defaults align to the tail of + # posonlyargs+args; kw_defaults align 1:1 with kwonlyargs (None = none). + _a = node.args + _defaulted = list( + zip( + (_a.posonlyargs + _a.args)[ + len(_a.posonlyargs) + len(_a.args) - len(_a.defaults) : + ], + _a.defaults, + ) + ) + [(p, d) for p, d in zip(_a.kwonlyargs, _a.kw_defaults) if d is not None] + for _param, _default in _defaulted: + if isinstance(_default, ast.Name): + _did = _default.id + if _did in open_aliases: + open_aliases.add(_param.arg) + elif _did in writer_aliases: + writer_aliases.add(_param.arg) + elif _did in archive_ctor_aliases: + archive_ctor_aliases.add(_param.arg) + elif _did in getattr_aliases: + getattr_aliases.add(_param.arg) + elif _did in partial_aliases: + partial_aliases.add(_param.arg) + elif _did in code_exec_aliases: + code_exec_aliases.add(_param.arg) + elif _did in dynamic_aliases: + dynamic_aliases.add(_param.arg) + elif isinstance(_default, ast.Attribute): + # An attribute writer / archive ctor / captured .open used as + # a default (def f(s=np.save), def f(z=zipfile.ZipFile), + # def f(o=Path('x').open)) binds the parameter like the + # equivalent assignment; a benign attribute (np.mean) does not. + if _default.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + writer_aliases.add(_param.arg) + elif _default.attr in _ARCHIVE_CTOR_NAMES: + archive_ctor_aliases.add(_param.arg) + elif _default.attr == "open": + dynamic_aliases.add(_param.arg) + elif ( + isinstance(_default, ast.Call) + and ( + ( + isinstance(_default.func, ast.Name) + and _default.func.id in partial_aliases + ) + or ( + isinstance(_default.func, ast.Attribute) + and _default.func.attr == "partial" + ) + ) + and _default.args + and _wraps_write_callable(_default.args[0]) + ): + dynamic_aliases.add(_param.arg) # def f(w=partial(open, mode="w")) + try: + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name.split(".")[0] in _AUTO_UNSAFE_PY_MODULES: + return True + elif isinstance(node, ast.ImportFrom): + if node.module and node.module.split(".")[0] in _AUTO_UNSAFE_PY_MODULES: + return True + # from-imports can bind mutating callables to bare names + # (from os import remove [as rm]); star imports hide anything. + for alias in node.names: + if alias.name == "*" or alias.name in _AUTO_UNSAFE_PY_ATTRS: + return True + # os.open imported as a bare callable is a low-level + # create/write, like the os.open attribute call below. + if alias.name == "open" and node.module in ("os", "posix"): + return True + elif isinstance(node, ast.Attribute): + # Any reference to a mutating attribute fails closed, even + # without an immediate call (rm = os.remove; rm("x")). + if node.attr in _AUTO_UNSAFE_PY_ATTRS: + return True + # builtins.exec / builtins.eval / builtins.__import__ (and + # compile/breakpoint) are dynamic code execution, matching the + # bare-name code_exec_aliases path; __builtins__.__import__(...) + # is a dynamic import that dodges the static import check. + if ( + node.attr in ("exec", "eval", "__import__", "breakpoint", "compile") + and isinstance(node.value, ast.Name) + and node.value.id in builtins_aliases + ): + return True + elif isinstance(node, ast.Name): + if node.id in code_exec_aliases: + return True + elif isinstance(node, ast.Constant): + # Credential paths / parent traversal in a string or bytes + # literal (open('/etc/passwd') and open(b'/etc/passwd')), or a + # glob that resolves to one (glob.glob('/e??/passwd')). + val = node.value + if isinstance(val, bytes): + val = val.decode("latin-1", "ignore") + if isinstance(val, str) and ( + _references_sensitive_path(val) or _glob_token_sensitive(val) + ): + return True + elif isinstance(node, (ast.BinOp, ast.JoinedStr)): + # A sensitive path concatenated from literals ('/etc'+'/passwd'), + # a pathlib / chain, an f-string (f'/proc/{pid}/environ'), a + # dynamic segment under a sensitive dir (f'/etc/{name}'), or one + # split through a literal variable (base = '/etc'; base+'/passwd'). + if _folded_is_sensitive( + _folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + ): + return True + elif isinstance(node, ast.Call): + # A sensitive path composed via os.path.join('/etc', name). + if _folded_is_sensitive( + _folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + ): + return True + func = node.func + # x.__call__(args) is just x(args): unwrap so open.__call__('o', + # 'w') / save.__call__(...) reach the open/writer checks below + # instead of looking like a harmless ".__call__" attribute call. + if isinstance(func, ast.Attribute) and func.attr == "__call__": + func = func.value + if isinstance(func, (ast.Call, ast.Subscript)): + return True # calling a call/subscript result is dynamic + # A concrete write callable (open/writer/archive-ctor alias, or a + # writer/mutating attribute) handed as an argument to any call + # escapes into a helper that can invoke it without a direct + # open()/writer site -- the same bypass the map/starmap/reduce + # branches below gate, but through a user-defined helper + # (def run(fn): fn('o','w').write('x'); run(open)). A benign + # callable argument (run(len)) is unaffected. + if any(_passed_write_callable(a) for a in node.args) or any( + _passed_write_callable(kw.value) for kw in node.keywords + ): + return True + if isinstance(func, ast.Name): + if func.id in dynamic_aliases: + return True # call through a getattr alias is dynamic + if func.id in open_aliases and _builtin_open_writes(node): + return True + # A writer imported as a bare name (from numpy import save). + if func.id in writer_aliases: + return True + # A bare archive constructor (from zipfile import ZipFile) + # takes the mode as its 2nd arg like open, so ZipFile(x, "w") + # writes but ZipFile(x) reads. + if func.id in archive_ctor_aliases and _builtin_open_writes(node): + return True + # A bare-imported logging.basicConfig(filename=...) opens a + # log file for writing (from logging import basicConfig). + if func.id in basicconfig_aliases and _basicconfig_writes(node): + return True + # A writer/open alias handed to a higher-order invoker + # (map(open, names, modes), starmap(np.save, ...), or an + # aliased m = map / sm = starmap) is called without a direct + # open(...)/save(...) site; the callable is the first + # positional arg. A benign map(len, ...) is unaffected. + if ( + func.id in invoker_aliases + and node.args + and _wraps_write_callable(node.args[0]) + ): + return True + elif isinstance(func, ast.Attribute): + # Writer methods persist to disk without open() (np.save, + # img.save, plt.savefig, df.to_csv, json.dump); ask before + # they mutate the workdir in auto mode. + if func.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + return True + # logging.basicConfig(filename=...) opens a log file for write. + if func.attr == "basicConfig" and _basicconfig_writes(node): + return True + # A qualified higher-order invoker (itertools.starmap(open, ...), + # functools.reduce(open, ...)) calls its first arg like the bare + # map/filter form; the writer-check on that arg keeps a benign + # itertools.starmap(len, ...) / df.map(transform) safe. + if ( + func.attr in _HIGHER_ORDER_INVOKERS + and node.args + and _wraps_write_callable(node.args[0]) + ): + return True + # fileinput.input(..., inplace=True) rewrites a file in place; + # the default fileinput.input(...) only reads, so gate inplace. + if ( + func.attr == "input" + and isinstance(func.value, ast.Name) + and func.value.id in fileinput_aliases + and _fileinput_inplace(node) + ): + return True + # os.open() always creates/writes a file descriptor + # (tracked through import aliases: import os as o; o.open()). + if ( + func.attr == "open" + and isinstance(func.value, ast.Name) + and func.value.id in os_aliases + ): + return True + # A pickle-backed loader (torch.load, joblib.load) can execute + # code embedded in the file it deserializes. + if ( + func.attr == "load" + and isinstance(func.value, ast.Name) + and func.value.id in load_module_aliases + ): + return True + if func.attr == "open" and _attr_open_writes(node): + return True + # An open bound onto an attribute (box.f = open; box.f('o','w')) + # writes on 'w'/'a'/'x' like the builtin, so gate the attr name. + if func.attr in attr_open_aliases and _builtin_open_writes(node): + return True + # ZipFile/TarFile/GzipFile/BZ2File/LZMAFile take the mode as + # the 2nd arg (like builtin open), so ZipFile(name, "w") writes + # but ZipFile(name) reads. + if func.attr in _ARCHIVE_CTOR_NAMES and _builtin_open_writes(node): + return True + # Enumerating a directory outside the sandbox reads host + # filenames (and enables reading their contents) the direct + # /etc/passwd checks would prompt for: Path('/etc').iterdir(), + # os.scandir('/etc'), os.listdir('/home'), os.walk('/'), + # Path('/home').glob('*'), glob.glob('/home/*'). Gate when the + # target dir folds to an absolute/tilde/sensitive path; a + # relative dir (Path('.').iterdir(), glob.glob('src/*')) stays + # safe, and an unresolved dynamic dir is left to other checks. + _enum_dir = None + if func.attr == "iterdir": + _enum_dir = func.value + elif func.attr in ("glob", "rglob", "iglob"): + # Path('/home').glob('*') enumerates the receiver dir; + # glob.glob('/home/*') enumerates the pattern's root dir. + _recv = _folded_path( + func.value, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) + if isinstance(_recv, str) and _recv not in ("", "\x00"): + _enum_dir = func.value + elif node.args: + _enum_dir = node.args[0] + elif ( + func.attr in ("scandir", "listdir", "walk") + and isinstance(func.value, ast.Name) + and func.value.id in os_aliases + and node.args + ): + _enum_dir = node.args[0] + if _enum_dir is not None: + _folded_dir = _folded_path( + _enum_dir, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) + if isinstance(_folded_dir, str) and ( + _folded_dir.startswith("/") + or _folded_dir.startswith("~") + or _folded_is_sensitive(_folded_dir) + ): + return True + except Exception: + return True # unexpected AST shape: fail closed + return False + + +# Cloud-metadata / link-local hosts (mirrors the sandbox SSRF blocklist): a +# read-named HTTP MCP tool pointed at one (fetch_url +# {"url": "http://169.254.169.254/..."}) reads instance credentials, so it asks. +_MCP_METADATA_HOST_RE = re.compile( + r"169\.254\.\d{1,3}\.\d{1,3}|" + r"100\.100\.100\.\d{1,3}|" + r"fd00:ec2::254|" + r"metadata\.google\.internal|" + r"metadata\.tencentyun\.com|" + r"://metadata(?=[:/])", + re.IGNORECASE, +) + + +def _mcp_arguments_reference_sensitive(arguments) -> bool: + """True if any string in an MCP call's arguments names a credential path, a + credential/secret environment variable (get_env {"name": "OPENAI_API_KEY"}), + or a cloud-metadata host (fetch_url {"url": "http://169.254.169.254/..."}).""" + + def walk(value) -> bool: + if isinstance(value, str): + return ( + _references_sensitive_path(value) + or bool(_AUTO_SENSITIVE_MCP_NOUN_RE.search(value)) + or bool(_MCP_METADATA_HOST_RE.search(value)) + ) + if isinstance(value, dict): + return any(walk(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(walk(v) for v in value) + return False + + return walk(arguments) + + +# DDL object types CREATE / DROP / ALTER share (DROP FUNCTION and ALTER INDEX +# mutate just like CREATE INDEX). +_SQL_DDL_OBJECTS = ( + r"table|database|schema|index|view|function|procedure|trigger|" + r"sequence|role|user|extension|type|domain|aggregate|policy" +) +# Modifiers between the DDL verb and object (CREATE OR REPLACE VIEW, DROP +# MATERIALIZED VIEW, CREATE UNIQUE INDEX). +_SQL_DDL_MODIFIERS = ( + r"(?:(?:or\s+replace|unique|temp|temporary|global|local|materialized|recursive)\s+)*" +) +# A SQL identifier (bare, "quoted", `quoted`, [bracketed]), optionally +# schema-qualified, so UPDATE "users"/public.users/ONLY .../[users] SET all hit. +_SQL_IDENT = r'(?:\w+|"(?:[^"]|"")*"|`(?:[^`]|``)*`|\[[^\]]+\])' +_SQL_UPDATE_TARGET = r"(?:only\s+)?" + _SQL_IDENT + r"(?:\s*\.\s*" + _SQL_IDENT + r")*" +# A read-named MCP tool (query_database, run_query) can still carry a mutating +# SQL statement; match DML/DDL as whole statements (DELETE FROM, DROP TABLE) so +# a natural-language query that merely contains the word "delete" stays safe. +_MCP_ARG_MUTATION_RE = re.compile( + r"\b(?:delete\s+from|" + r"drop\s+" + _SQL_DDL_MODIFIERS + r"(?:" + _SQL_DDL_OBJECTS + r")|" + # Match the whole identifier (the outer trailing \b needs the alternative to + # end on a word boundary, so a bare \w stops mid-name and TRUNCATE users slips + # through); the optional opening quote/bracket/backtick covers "users"/[users]. + r"truncate\s+(?:table\s+)?[\"\[`]?\w+|" + # UPDATE [AS alias] SET: allow an explicit AS alias before SET so + # UPDATE users AS u SET is caught, not just the bare form. The implicit-alias + # form (UPDATE users u SET) is left out because it is indistinguishable from + # the prose "update set" and would flag natural language. + r"update\s+" + _SQL_UPDATE_TARGET + r"(?:\s+as\s+" + _SQL_IDENT + r")?\s+set\b|" + r"insert\s+into|replace\s+into|" + # SELECT ... INTO OUTFILE/DUMPFILE writes a file (MySQL); bare SELECT INTO + # is left out (PL/pgSQL uses it to read into a variable). + r"select\s+[^;]*?\binto\s+(?:outfile|dumpfile)\b|" + # ALTER SYSTEM persists PostgreSQL server configuration; SYSTEM is not one of + # the DDL objects above, so match it explicitly. + r"alter\s+system\b|" + r"alter\s+" + _SQL_DDL_MODIFIERS + r"(?:" + _SQL_DDL_OBJECTS + r")|" + r"create\s+" + _SQL_DDL_MODIFIERS + r"(?:" + _SQL_DDL_OBJECTS + r")|" + r"grant\s+\w+|revoke\s+\w+|merge\s+into|" + # Catalog mutations: COMMENT ON , SECURITY LABEL, and LOCK TABLE change + # metadata or take a lock. Each needs a following keyword, so a "comment" + # column (SELECT comment FROM t) or "locks" table stays safe. + r"comment\s+on\b|security\s+label\b|lock\s+table\b|" + # PostgreSQL maintenance writes: REFRESH MATERIALIZED VIEW rewrites the view, + # REINDEX rebuilds an index. Both need a following object keyword/name, so a + # column or word "refresh"/"reindex" in prose stays safe. + r"refresh\s+materialized\s+view|reindex\s+\w+|" + # CALL proc(...) / EXEC[UTE] name / VACUUM mutate; CALL needs a following + # "(", ";", or end so natural-language "call me back" stays safe. + r"call\s+\w+(?=\s*[(;]|\s*$)|exec(?:ute)?\s+\w+|vacuum|" + # COPY ... FROM bulk-loads and COPY ... TO writes a file ([^;] stays in one + # statement). + r"copy\s+[^;]*?\b(?:from|to)\b)\b", + re.IGNORECASE, +) +# SQLite statements the base regex misses: ATTACH/DETACH a database (DATABASE +# optional via the quoted-path form), a write-form PRAGMA (name=value / name(...), +# unlike the read-form PRAGMA name), and load_extension() which runs a shared +# library. These tokens are not natural language, so benign text does not trip. +_MCP_ARG_SQLITE_MUTATION_RE = re.compile( + r"\b(?:attach|detach)\s+database\b" + r"|\battach\s+(?:database\s+)?['\"]" + r"|\bpragma\s+\w+(?:\.\w+)?\s*(?:=|\()" + r"|\bload_extension\s*\(", + re.IGNORECASE, +) +# State-changing SQL functions that mutate or write files inside a read-shaped +# SELECT (pg_terminate_backend, setval, pg_write_file, lo_export, ...). The +# trailing "(" is required, so a column named setval_count stays safe. +_MCP_ARG_SQL_FUNCTION_RE = re.compile( + r"\b(?:pg_terminate_backend|pg_cancel_backend|pg_write_file|lo_export|" + r"lo_import|setval|nextval|set_config|pg_notify|dblink_exec|pg_reload_conf|" + r"pg_rotate_logfile|" + # advisory locks change session/transaction lock state (read-shaped SELECT). + r"pg_advisory_(?:lock|lock_shared|unlock|unlock_shared|unlock_all|" + r"xact_lock|xact_lock_shared)|" + r"pg_try_advisory_(?:lock|lock_shared|xact_lock|xact_lock_shared))\s*\(", + re.IGNORECASE, +) +# SQL engines treat /* */ and -- comments as whitespace, so DELETE/**/FROM and +# UPDATE/**/users evade the \s+ in the mutation regex; collapse comments to a +# space before matching. +_SQL_COMMENT_RE = re.compile(r"/\*.*?\*/|--[^\n]*", re.DOTALL) +# A GraphQL mutation on a read-named tool. Directives are valid between the name +# and body (mutation M @audit { ... }), so allow @directive[(args)] before ( or {. +_GRAPHQL_MUTATION_RE = re.compile( + r"\bmutation\b\s*\w*\s*(?:@\w+(?:\s*\([^)]*\))?\s*)*[({]", re.IGNORECASE +) +# GraphQL # comments run to end-of-line and count as whitespace, so a comment +# between `mutation` and the body (mutation # note\n { ... }) would otherwise +# hide it; collapse them to a space before matching. +_GRAPHQL_COMMENT_RE = re.compile(r"#[^\n]*") + + +# HTTP verbs that mutate the target resource; a generic HTTP MCP tool +# (mcp__http__get_url {"method": "DELETE"}) mutates an external service even +# though its name looks read-only. GET/HEAD/OPTIONS/TRACE only read. +_MUTATING_HTTP_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"}) +_HTTP_METHOD_KEYS = frozenset({"method", "http_method", "httpmethod", "verb", "http_verb"}) + + +def _mcp_arguments_mutate(arguments) -> bool: + """True if an MCP call's arguments carry a mutating command, so a read-named + but write-capable tool (query_database {"query": "DELETE FROM runs"}, + query_graphql {"query": "mutation { deleteIssue(id: 1) }"}, or an HTTP tool + {"method": "DELETE"}) asks.""" + + def walk(value) -> bool: + if isinstance(value, str): + _sql = _SQL_COMMENT_RE.sub(" ", value) + return ( + bool(_MCP_ARG_MUTATION_RE.search(_sql)) + or bool(_MCP_ARG_SQLITE_MUTATION_RE.search(_sql)) + or bool(_MCP_ARG_SQL_FUNCTION_RE.search(_sql)) + or bool(_GRAPHQL_MUTATION_RE.search(_GRAPHQL_COMMENT_RE.sub(" ", value))) + ) + if isinstance(value, dict): + for k, v in value.items(): + if ( + isinstance(k, str) + and k.lower() in _HTTP_METHOD_KEYS + and isinstance(v, str) + and v.strip().upper() in _MUTATING_HTTP_METHODS + ): + return True + return any(walk(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(walk(v) for v in value) + return False + + return walk(arguments) + + +# Tools that are read-only / non state-mutating regardless of their arguments, +# so auto mode never has to pause them (their safety needs no argument scan). +# render_html is NOT unconditionally safe: it runs arbitrary HTML/JS in the +# canvas preview frame. A static canvas (charts, layout, inline SVG) never +# reaches the network, but code that calls out can exfiltrate or fetch under the +# preview's CSP when artifact network access is enabled, so those ask; a canvas +# with no network construct still auto-runs. Matches JS egress APIs, a remote or +# root-relative ") is False + ) + assert rh("") is False + assert rh("") is False + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + # Worker / SharedWorker constructors run an off-thread script the scan cannot + # see (a module worker from a CORS CDN, or a blob/same-origin worker that + # fetches/importScripts) under worker-src http: https: blob:, so they ask. + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is False # not a ctor + assert rh("") is False # unrelated class, not a real Worker + # Resource-loading forms beyond a direct fetch also reach the network. + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True # root-relative resolves to origin + assert rh("") is True # protocol-relative + # Self-navigation sinks exfiltrate by navigating the frame away. + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is False # reload is not navigation + assert rh("") is False + # Obfuscated egress: a block comment splitting fetch(, or bracket access. + assert rh("") is True + assert rh("") is True + # A computed bracket key spliced from string fragments on a global host object. + assert rh("") is True + assert rh("") is True + # A computed key on a plain object (not a global host) stays a static canvas. + assert rh("") is False + assert rh("") is False # comment only + # A meta-refresh with a url navigates the frame to an external origin. + assert rh('') is True + assert rh("") is True + assert rh('') is False # self-reload, no url + assert rh('

Hi

') is False # ordinary meta stays safe + + +def test_unknown_tools_fail_closed(): + assert is_potentially_unsafe_tool_call("mystery_tool", {}) is True + + +def test_is_always_safe_tool(): + from core.inference.tools import is_always_safe_tool + for name in ("web_search", "search_knowledge_base"): + assert is_always_safe_tool(name) is True + # render_html is no longer unconditionally safe: a networked canvas can prompt, + # which cannot be judged before its arguments stream. + for name in ("python", "terminal", "mystery_tool", "mcp__srv__read", "render_html"): + assert is_always_safe_tool(name) is False + + +@pytest.mark.parametrize( + ("tool", "unsafe"), + [ + ("get_weather", False), + ("list_files", False), + ("search", False), + ("send_email", True), + ("create_issue", True), + ("delete_row", True), + ("get_or_create_issue", True), # mutating verb overrides read prefix + ("read_and_delete_file", True), + ("find_and_update_row", True), + ("get_and_commit_changes", True), # commit/save/archive are mutating + ("read_and_save_file", True), + ("list_and_archive", True), + ("list_and_clone_repo", True), # clone/checkout/comment are mutating + ("fetch_and_comment_issue", True), + ("get_and_checkout_branch", True), + ("read_and_append_file", True), # append/prepend are mutating + ("prepend_line", True), + ("get_and_upsert_row", True), # upsert/assign are mutating + ("list_and_assign_issue", True), + ("read_and_copy_file", True), # copy-style verbs create/overwrite state + ("get_and_copy_resource", True), + ("read_and_duplicate_entry", True), + ("fetch_and_download_asset", True), # download writes local state + ("list_and_export_data", True), # import/export/backup/restore/snapshot + ("get_and_snapshot_volume", True), + ("get_and_mark_read", True), # mark/subscribe change external state + ("get_and_subscribe", True), + ("list_and_unsubscribe", True), + ("get_and_reply_email", True), # reply/notify send/change external state + ("list_and_notify_users", True), + ("read_secret", True), # credential noun: a read that discloses a secret + ("list_tokens", True), + ("get_credentials", True), + ("fetch_api_key", True), # scoped *_key noun + ("read_access_key", True), + ("get_password", True), + ("read_passphrase", True), + ("read_report", False), # plain read stays safe + ("get_primary_key", False), # a schema key is not a credential + ("search_keyboard_shortcuts", False), # 'key' inside another word stays safe + ("list_bookmarks", False), # 'mark' substring in a token stays safe + ("list_notifications", False), # 'notify' is a different token than 'notifications' + ], +) +def test_mcp_classifier(tool, unsafe): + name = f"{MCP_TOOL_PREFIX}srv1__{tool}" + assert is_potentially_unsafe_tool_call(name, {}) is unsafe + + +@pytest.mark.parametrize( + ("args", "unsafe"), + [ + ({"path": "/etc/passwd"}, True), # read-named tool at a credential path + ({"path": "../../.ssh/id_rsa"}, True), + ({"nested": {"file": "~/.aws/credentials"}}, True), + ({"name": "OPENAI_API_KEY"}, True), # explicit credential env-var read + ({"name": "AWS_SECRET_ACCESS_KEY"}, True), + ({"key": "DATABASE_PASSWORD"}, True), + ( + {"url": "http://169.254.169.254/latest/meta-data/iam/security-credentials/"}, + True, + ), # AWS instance-metadata host + ( + {"url": "http://metadata.google.internal/computeMetadata/v1/"}, + True, + ), # GCP metadata host + ({"path": "notes.txt"}, False), # ordinary path stays safe + ({"path": "data/report.csv"}, False), + ({"name": "PATH"}, False), # a non-secret env var stays safe + ({"name": "HOME"}, False), + ({"url": "https://example.com/api"}, False), # ordinary URL stays safe + ({"url": "http://localhost:8080/health"}, False), # localhost app stays safe + ], +) +def test_mcp_sensitive_arguments(args, unsafe): + name = f"{MCP_TOOL_PREFIX}fs__read_file" + assert is_potentially_unsafe_tool_call(name, args) is unsafe + + +@pytest.mark.parametrize( + ("args", "unsafe"), + [ + ({"query": "DELETE FROM runs"}, True), # read-named tool, mutating query + ({"sql": "DROP TABLE users"}, True), + ({"query": "UPDATE t SET x=1"}, True), + ({"query": "INSERT INTO t VALUES (1)"}, True), + ({"query": "SELECT * FROM runs"}, False), # read query stays safe + ({"query": "how to delete old files"}, False), # NL text with 'delete' stays safe + ({"query": "find the created_at column"}, False), # 'created' substring stays safe + ({"query": "DELETE/**/FROM runs"}, True), # inline SQL comment as whitespace + ({"query": "UPDATE/**/t SET x=1"}, True), + ({"query": "DROP/**/TABLE users"}, True), + ({"query": "SELECT * FROM runs -- delete later"}, False), # trailing comment stays safe + ({"query": "COPY users FROM '/tmp/u.csv'"}, True), # bulk load writes the table + ({"query": "COPY users (id, name)\nFROM STDIN"}, True), # multiline COPY FROM + ({"query": "COPY (SELECT 1) TO '/tmp/o.csv'"}, True), # COPY TO writes a server file + ({"query": "SELECT copy_count FROM t"}, False), # 'copy' substring column stays safe + ({"query": "mutation { deleteIssue(id: 1) }"}, True), # GraphQL mutation + ({"query": "mutation DelIssue { deleteIssue(id: 1) }"}, True), # named GraphQL mutation + ({"query": "mutation # note\n { deleteIssue(id: 1) }"}, True), # comment before body + ({"query": "mutation # c\n Del { deleteIssue(id: 1) }"}, True), # comment before name + ({"query": "query { issue(id: 1) { title } }"}, False), # GraphQL read query stays safe + ({"query": "{ issue(id: 1) { title } }"}, False), # shorthand GraphQL query stays safe + ({"query": "query # note\n { issue(id: 1) }"}, False), # commented read query stays safe + ({"query": "CREATE OR REPLACE VIEW v AS SELECT 1"}, True), # DDL with a modifier + ({"query": "CREATE UNIQUE INDEX idx ON t(x)"}, True), # DDL with UNIQUE + ({"query": "CREATE TEMP TABLE t (id int)"}, True), # DDL with TEMP + ({"query": "CREATE MATERIALIZED VIEW mv AS SELECT 1"}, True), # materialized view DDL + ({"query": "CREATE FUNCTION f() RETURNS int AS $$ $$"}, True), # function DDL + ({"query": "ALTER SYSTEM SET work_mem = '1GB'"}, True), # persists server config + ({"query": "alter system reset all"}, True), # ALTER SYSTEM RESET + ({"query": "SELECT * FROM system_logs"}, False), # 'system' as a table name stays safe + ({"query": "SELECT * FROM created_view"}, False), # 'create' substring stays safe + ({"query": "CALL delete_all_users()"}, True), # stored procedure invocation + ({"query": "EXEC purge_queue"}, True), # EXEC procedure + ({"query": "EXECUTE sp_drop"}, True), # EXECUTE procedure + ({"query": "VACUUM INTO 'backup.db'"}, True), # VACUUM rewrites the database + ({"query": "please call me back later"}, False), # NL 'call' stays safe + ({"query": "ATTACH DATABASE '/tmp/x.db' AS x"}, True), # attaches a database file + ({"query": "DETACH DATABASE x"}, True), # detaches a database + ({"query": "PRAGMA user_version = 42"}, True), # write-form PRAGMA + ({"query": "PRAGMA journal_mode=WAL"}, True), # write-form PRAGMA (no spaces) + ({"query": "PRAGMA foreign_keys(0)"}, True), # call-form PRAGMA write + ({"query": "SELECT load_extension('/tmp/evil.so')"}, True), # loads native code + ({"query": "PRAGMA journal_mode"}, False), # read-form PRAGMA stays safe + ({"query": "can you attach the report to the email"}, False), # NL 'attach' stays safe + ({"query": "ATTACH '/tmp/x.db' AS x"}, True), # ATTACH without DATABASE keyword + ({"query": "PRAGMA main.user_version = 1"}, True), # schema-qualified write PRAGMA + ({"query": "attach it as draft"}, False), # NL 'attach ... as' stays safe + ({"query": "DROP FUNCTION f()"}, True), # DROP of a non-table object + ({"query": "ALTER INDEX idx RENAME TO idx2"}, True), # ALTER of a non-table object + ({"query": "DROP MATERIALIZED VIEW mv"}, True), # DROP with a modifier + ({"query": "ALTER USER bob WITH PASSWORD 'x'"}, True), # ALTER USER mutates + ({"query": "SELECT dropped_at FROM t"}, False), # 'drop' substring column stays safe + ({"query": "mutation M @audit { deleteIssue(id: 1) }"}, True), # directive GraphQL mutation + ( + {"query": "query Q @cached { issue(id: 1) { title } }"}, + False, + ), # directive GraphQL read stays safe + ({"query": 'UPDATE "users" SET admin=1'}, True), # double-quoted UPDATE target + ({"query": "UPDATE public.users SET admin=1"}, True), # schema-qualified UPDATE + ({"query": "UPDATE ONLY public.users SET admin=1"}, True), # ONLY-qualified UPDATE + ({"query": "UPDATE `users` SET admin=1"}, True), # backtick-quoted UPDATE + ({"query": "UPDATE [users] SET admin=1"}, True), # bracket-quoted UPDATE + ({"query": "please update the documentation set"}, False), # NL 'update ... set' stays safe + ({"query": "SELECT pg_terminate_backend(123)"}, True), # state-changing SQL function + ({"query": "SELECT setval('s', 1)"}, True), # sequence mutation function + ({"query": "SELECT pg_write_file('/tmp/p', 'x')"}, True), # server-side file write + ({"query": "SELECT lo_export(123, '/tmp/p')"}, True), # large-object export to a file + ({"query": "SELECT setval_col FROM t"}, False), # 'setval' column prefix stays safe + ( + {"query": "SELECT secret INTO OUTFILE '/tmp/leak' FROM users"}, + True, + ), # INTO OUTFILE write + ({"query": "SELECT x INTO DUMPFILE '/tmp/d' FROM t"}, True), # INTO DUMPFILE write + ( + {"query": "SELECT count(*) INTO cnt FROM t"}, + False, + ), # PL/pgSQL SELECT INTO var stays safe + ({"query": "REFRESH MATERIALIZED VIEW mv"}, True), # materialized view rewrite + ({"query": "REINDEX INDEX idx"}, True), # index rebuild + ({"query": "REINDEX TABLE t"}, True), # table reindex + ({"query": "SELECT refresh_count FROM t"}, False), # 'refresh' column stays safe + ({"query": "please refresh the page"}, False), # NL 'refresh' stays safe + ({"query": "COMMENT ON TABLE users IS 'owned'"}, True), # catalog metadata write + ({"query": "LOCK TABLE users IN ACCESS EXCLUSIVE MODE"}, True), # explicit lock + ({"query": "SECURITY LABEL FOR x ON TABLE t IS 'z'"}, True), # security label write + ({"query": "CREATE POLICY p ON accounts USING (true)"}, True), # row-security policy DDL + ({"query": "SELECT comment FROM t"}, False), # 'comment' column stays safe + ({"query": "SELECT * FROM locks"}, False), # 'locks' table stays safe + ({"query": "SELECT nextval('billing_seq')"}, True), # sequence advance mutates + ({"query": "SELECT pg_advisory_lock(42)"}, True), # advisory lock changes state + ({"query": "SELECT pg_notify('jobs', 'wake')"}, True), # server-side notification + ({"query": "SELECT set_config('x', 'y', false)"}, True), # session config write + ({"query": "SELECT nextval_col FROM t"}, False), # 'nextval' column prefix stays safe + ({"query": "TRUNCATE users"}, True), # multi-char table name (bare TRUNCATE) + ({"query": "TRUNCATE TABLE accounts"}, True), # multi-char TRUNCATE TABLE + ({"query": 'TRUNCATE TABLE "users"'}, True), # quoted TRUNCATE target + ({"query": "TRUNCATE accounts RESTART IDENTITY"}, True), # TRUNCATE with options + ({"query": "SELECT truncate_log FROM t"}, False), # 'truncate' column stays safe + ({"query": "UPDATE users AS u SET admin=1"}, True), # aliased UPDATE target (AS) + ({"query": 'UPDATE "users" AS u SET x=1'}, True), # quoted+aliased UPDATE + ({"query": "UPDATE public.users AS u SET x=1"}, True), # schema-qualified aliased UPDATE + ({"query": "SELECT * FROM users AS u"}, False), # aliased SELECT stays safe + ({"query": "please update the documentation set"}, False), # NL, no AS, stays safe + ({"query": "GRANT SELECT ON t TO u"}, True), # privilege grant (multi-word) + ({"query": "REVOKE ALL ON t FROM u"}, True), # privilege revoke (multi-word) + ({"query": "SELECT * FROM grants"}, False), # 'grants' table stays safe + ({"url": "http://x", "method": "DELETE"}, True), # mutating HTTP verb arg + ({"method": "POST"}, True), + ({"verb": "PUT"}, True), # alternate method-key name + ({"method": "GET"}, False), # read HTTP verb stays safe + ({"method": "HEAD"}, False), + ], +) +def test_mcp_mutating_arguments(args, unsafe): + name = f"{MCP_TOOL_PREFIX}db__query_database" + assert is_potentially_unsafe_tool_call(name, args) is unsafe + + +# ── loop behavior ─────────────────────────────────────────────────── + +_DEFAULT_TOOLS = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "web_search"}}, +] + + +class _FakeExecuteTool: + def __init__(self): + self.calls = [] + self.disable_sandbox_seen = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + thread_id = None, + rag_scope = None, + disable_sandbox = False, + ): + self.calls.append((name, arguments)) + self.disable_sandbox_seen.append(disable_sandbox) + return f"RESULT[{name}]" + + +def _tool_call(name, args_json): + return f'{{"name": "{name}", "arguments": {args_json}}}' + + +def _multi_turn(turns): + turn_iter = iter(turns) + + def _gen(_messages): + try: + yield next(turn_iter) + except StopIteration: + return + + return _gen + + +def _drive(turns, decisions, **loop_kwargs): + """Run the loop, resolving each gated tool_start with the next decision.""" + decision_iter = iter(decisions) + exec_fn = _FakeExecuteTool() + # A per-call session id so a leaked pending approval from another test can + # never collide with this run's approval registry entries. + session = f"{_SESSION}-{uuid.uuid4().hex}" + gen = run_safetensors_tool_loop( + single_turn = _multi_turn(turns), + messages = [{"role": "user", "content": "hi"}], + tools = _DEFAULT_TOOLS, + execute_tool = exec_fn, + session_id = session, + **loop_kwargs, + ) + events = [] + for ev in gen: + events.append(ev) + if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): + resolve_tool_decision(ev["approval_id"], next(decision_iter), session_id = session) + return events, exec_fn + + +def _tool_starts(events): + return [e for e in events if e["type"] == "tool_start"] + + +def _diag(events, exec_fn): + """A compact dump of what the loop actually did, attached to the loop-driving + assertions so a full-suite-only failure on CI (which does not reproduce when + the file runs alone) reports the real event stream instead of a bare diff.""" + return ( + f"calls={exec_fn.calls} sandbox_seen={exec_fn.disable_sandbox_seen} " + f"events={[(e.get('type'), e.get('awaiting_confirmation'), e.get('tool_name')) for e in events]}" + ) + + +def test_auto_mode_does_not_gate_safe_calls(): + events, exec_fn = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + [], + confirm_tool_calls = True, + permission_mode = "auto", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts[0]["approval_id"] == "" + assert exec_fn.calls == [("python", {"code": "print(1)"})], _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [False], _diag( + events, exec_fn + ) # sandbox stays on in auto + + +def test_auto_mode_gates_unsafe_calls(): + events, exec_fn = _drive( + [_tool_call("python", '{"code": "import os; os.remove(\\"x\\")"}'), "final"], + ["allow"], + confirm_tool_calls = True, + permission_mode = "auto", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is True, _diag(events, exec_fn) + assert starts[0]["approval_id"] + assert len(exec_fn.calls) == 1, _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [False], _diag(events, exec_fn) + + +def test_ask_mode_gates_even_safe_calls(): + events, _ = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + ["allow"], + confirm_tool_calls = True, + permission_mode = "ask", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is True + + +def test_unset_mode_behaves_as_ask(): + events, _ = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + ["allow"], + confirm_tool_calls = True, + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is True + + +def test_off_mode_never_gates_and_keeps_sandbox(): + # "Off": no prompts even for unsafe calls, but the sandbox stays on. + events, exec_fn = _drive( + [_tool_call("python", '{"code": "import os; os.remove(\\"x\\")"}'), "final"], + [], + confirm_tool_calls = True, # off must win over a stray confirm flag + permission_mode = "off", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts[0]["approval_id"] == "" + assert exec_fn.disable_sandbox_seen == [False], _diag(events, exec_fn) + + +def test_full_mode_never_gates_and_drops_sandbox(): + events, exec_fn = _drive( + [_tool_call("python", '{"code": "import os; os.remove(\\"x\\")"}'), "final"], + [], + confirm_tool_calls = True, # full must win over the confirm gate + permission_mode = "full", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [True], _diag(events, exec_fn) + + +def test_bypass_flag_implies_full_mode(): + # Legacy callers that only set bypass_permissions keep the same behavior. + events, exec_fn = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + [], + confirm_tool_calls = True, + bypass_permissions = True, + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [True], _diag(events, exec_fn) + + +def test_bypass_permissions_folds_to_full_on_request_models(): + # A legacy bypass caller that also sends a stale ask/auto mode normalizes to + # full, so the route guards (which reject ask/auto) don't 400 the request. + for cls in (ChatCompletionRequest, AnthropicMessagesRequest): + req = cls( + messages = [{"role": "user", "content": "hi"}], + bypass_permissions = True, + permission_mode = "auto", + ) + assert req.permission_mode == "full" + assert req.bypass_permissions is True + + +def test_unknown_permission_mode_normalizes_to_ask_on_request_models(): + # An unrecognized mode from a newer UI/client must degrade to the safest gate + # ("ask") at the API boundary instead of a 422, so the forward-compat fallback + # the tool loops already apply (unknown -> ask) is reachable. None stays unset; + # the four known modes pass through untouched. + for cls in (ChatCompletionRequest, AnthropicMessagesRequest): + for unknown in ("paranoid", "readonly", "bogus", ""): + req = cls( + messages = [{"role": "user", "content": "hi"}], + permission_mode = unknown, + ) + assert req.permission_mode == "ask", (cls.__name__, unknown) + assert ( + cls(messages = [{"role": "user", "content": "hi"}], permission_mode = None).permission_mode + is None + ) + for known in ("ask", "auto", "off", "full"): + req = cls( + messages = [{"role": "user", "content": "hi"}], + permission_mode = known, + ) + # 'full' folds to bypass but the mode string is preserved. + assert req.permission_mode == known, (cls.__name__, known) + + +def test_ask_auto_self_enable_confirm_on_chat_request(): + # "Ask" gates every call, so a direct /chat/completions caller that requests + # ask but omits the legacy confirm flag self-enables it when Studio's own tool + # loop is requested. Only the router's loop-entry signals count (enable_tools / + # mcp_enabled); enabled_tools alone never starts the loop. + for loop in ({"enable_tools": True}, {"mcp_enabled": True}): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + **loop, + ) + assert req.confirm_tool_calls is True + # "auto" is NOT folded: it only prompts for a classifier-flagged call, so + # leaving confirm unset lets the route apply the safe-only-selection exception + # (a safe-only auto request needs no stream) instead of an explicit confirm + # forcing stream=true. The mode still drives the loop's per-call gate. + for loop in ({"enable_tools": True}, {"mcp_enabled": True}): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "auto", + **loop, + ) + assert req.confirm_tool_calls is None + # enabled_tools by itself is a passthrough filter, not a loop-entry signal: + # a client-tool passthrough that also lists enabled_tools must route verbatim + # (confirm stays unset), else the confirm-without-stream guard 400s it. + for mode in ("ask", "auto"): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = mode, + enabled_tools = ["terminal"], + tools = [{"type": "function", "function": {"name": "f"}}], + ) + assert req.confirm_tool_calls is None + # An explicit confirm_tool_calls=False wins over the ask mode (opts out of the + # gate), matching _permission_mode_confirm and the Anthropic pre-switch guard; + # the fold only self-enables when the flag is unset, so a caller cannot get a + # different answer on the chat path than the Anthropic path for the same body. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + enable_tools = True, + confirm_tool_calls = False, + ) + assert req.confirm_tool_calls is False + # A plain client-tool passthrough (client-supplied tools that Studio does not + # execute) must NOT self-enable confirm, or the route rejects the passthrough. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + tools = [{"type": "function", "function": {"name": "f"}}], + ) + assert req.confirm_tool_calls is None + # ask/auto without any tool request has nothing to gate; confirm stays unset. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + ) + assert req.confirm_tool_calls is None + # Legacy callers with no permission_mode keep their confirm flag untouched. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + confirm_tool_calls = False, + ) + assert req.confirm_tool_calls is False + # External-provider requests are not folded (the provider branch rejects + # confirm_tool_calls with tools, and permission_mode is a local concept). + for extra in ({"provider_id": "p1"}, {"provider_type": "openai"}): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + enable_tools = True, + **extra, + ) + assert req.confirm_tool_calls is None + + +def test_permission_mode_confirm_derivation(): + # The route derives the effective confirm gate from permission_mode so that a + # tool loop forced on by CLI policy (no request-level tool flag) still honors + # the documented "unset behaves as ask" default. + from routes.inference import _permission_mode_confirm + + def req(**kw): + return ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}], **kw) + + # An explicit confirm flag always wins (True gates, False opts out). + assert _permission_mode_confirm(req(confirm_tool_calls = True, stream = False)) is True + assert _permission_mode_confirm(req(confirm_tool_calls = False, permission_mode = "ask")) is False + # Explicit ask/auto always engage the gate (a non-streaming one is rejected + # by the guard that reads this). + assert _permission_mode_confirm(req(permission_mode = "ask", stream = False)) is True + assert _permission_mode_confirm(req(permission_mode = "auto", stream = False)) is True + # off/full never prompt. + assert _permission_mode_confirm(req(permission_mode = "off")) is False + assert _permission_mode_confirm(req(permission_mode = "full")) is False + # An unset mode defaults to ask, but only realizably on a streaming request; + # a non-streaming unset request keeps the legacy run-without-gate behavior. + assert _permission_mode_confirm(req(stream = True)) is True + assert _permission_mode_confirm(req(stream = False)) is False + + +def test_confirm_gate_needs_stream(): + # auto only prompts for a classifier-flagged call, so an auto request that can + # only select always-safe tools (web_search / RAG) needs no stream and must not + # be rejected by the confirm-without-stream guard. + from routes.inference import _confirm_gate_needs_stream + + def req(**kw): + return ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}], **kw) + + safe = ["web_search", "search_knowledge_base"] + # auto + a safe-only selection never prompts -> no stream needed. + assert _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = safe)) is False + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = ["web_search"])) + is False + ) + # render_html can prompt when its canvas reaches the network, so a selection + # that includes it needs a stream to deliver that prompt. + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search", "render_html"]) + ) + is True + ) + # But a selectable unsafe tool, an unrestricted (omitted) selection, MCP, or an + # explicit confirm flag all still require streaming under auto. + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = ["terminal"])) is True + ) + assert _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True)) is True + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search"], mcp_enabled = True) + ) + is True + ) + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search"], confirm_tool_calls = True) + ) + is True + ) + # An explicit empty selection runs no built-in tool, so nothing can prompt and + # no stream is needed (distinct from an omitted list, which means all tools). + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True, enabled_tools = [])) + is False + ) + # ask prompts for every call, so even a safe-only selection needs streaming. + assert _confirm_gate_needs_stream(req(permission_mode = "ask", enabled_tools = safe)) is True + # off/full never prompt; unset non-streaming keeps the legacy run-without-gate. + assert _confirm_gate_needs_stream(req(permission_mode = "off", enabled_tools = safe)) is False + assert _confirm_gate_needs_stream(req(permission_mode = "full", enabled_tools = safe)) is False + assert _confirm_gate_needs_stream(req(enabled_tools = safe, stream = False)) is False diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 63fdbbd8e9..eae1a75161 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2592,6 +2592,50 @@ class TestLoopBasic: assert tool_starts[0]["arguments"] == {} assert "" 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( + [ + [ + "", + "", + "Hi", + ], + ["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 "" 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( [ diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index 1f7608a4fc..2c13e13bbb 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -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(): diff --git a/studio/backend/tests/test_shutdown_preserves_live_worker.py b/studio/backend/tests/test_shutdown_preserves_live_worker.py new file mode 100644 index 0000000000..faf273411c --- /dev/null +++ b/studio/backend/tests/test_shutdown_preserves_live_worker.py @@ -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 diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py new file mode 100644 index 0000000000..457dfc8ea2 --- /dev/null +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -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" diff --git a/studio/backend/tests/test_transformers_latest.py b/studio/backend/tests/test_transformers_latest.py new file mode 100644 index 0000000000..20616dccba --- /dev/null +++ b/studio/backend/tests/test_transformers_latest.py @@ -0,0 +1,1099 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the latest-transformers support check and the consented sidecar install.""" + +import ast +import json +import os +import textwrap +import time +import pytest +from pathlib import Path + + +# The backend uses "from utils..." imports; ensure the backend dir is on sys.path. +import sys + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the custom logger before importing the modules under test. +import types as _types + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import utils.transformers_latest as tl +import utils.transformers_version as tv +from utils.transformers_latest import ( + check_upgrade_for_model, + install_latest_transformers, + latest_transformers_supports, + _fetch_remote_model_types, + _model_types_from_config, +) +from utils.transformers_version import ( + _config_mapping_cache, + _config_json_cache, + _higher_tier, + _is_valid_version_string, + _model_types_from_source, + _tier_from_config_mapping, + _venv_t5_latest_packages, + activate_transformers_for_subprocess, + ensure_latest_transformers_venv, + get_transformers_tier, + latest_venv_pinned_version, +) + + +# A CONFIG_MAPPING_NAMES source exercising every construct the AST extractor supports. +_MAPPING_SOURCE = """ +from collections import OrderedDict +CONFIG_MAPPING_NAMES = OrderedDict( + [ + ("llama", "LlamaConfig"), + ("gemma4", "Gemma4Config"), + ], + **{"qwen3_moe": "Qwen3MoeConfig"}, +) +CONFIG_MAPPING_NAMES.update({"brandnew_arch": "BrandNewConfig"}) +""" + +_MAIN_ONLY_SOURCE = """ +CONFIG_MAPPING_NAMES = { + "llama": "LlamaConfig", + "gemma4": "Gemma4Config", + "qwen3_moe": "Qwen3MoeConfig", + "brandnew_arch": "BrandNewConfig", + "dev_only_arch": "DevOnlyConfig", +} +""" + + +class _FakeResponse: + def __init__(self, body: bytes): + self._body = body + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +def _fake_urlopen_factory(counter: dict): + """urlopen stub serving the PyPI JSON and both refs' mapping sources.""" + + def _fake_urlopen(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + counter[url] = counter.get(url, 0) + 1 + counter["__total__"] = counter.get("__total__", 0) + 1 + if url == tl._PYPI_JSON_URL: + return _FakeResponse(json.dumps({"info": {"version": "5.13.0"}}).encode()) + if "/v5.13.0/" in url and url.endswith("auto_mappings.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + if "/v5.13.0/" in url and url.endswith("configuration_auto.py"): + return _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n") + if "/main/" in url and url.endswith("auto_mappings.py"): + return _FakeResponse(_MAIN_ONLY_SOURCE.encode()) + if "/main/" in url and url.endswith("configuration_auto.py"): + return _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n") + raise AssertionError(f"unexpected URL fetched: {url}") + + return _fake_urlopen + + +@pytest.fixture(autouse = True) +def _isolated_caches(tmp_path: Path, monkeypatch): + """Fresh in-memory + on-disk caches per test; no accidental real studio_root writes.""" + tl.clear_caches() + monkeypatch.setattr(tl, "_cache_file", lambda: tmp_path / "transformers_latest_check.json") + # The sidecar swap reservation writes a lock file next to the venv dir; + # point it at tmp so tests never touch the real studio root. + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", raising = False) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + tl.clear_caches() + + +def _no_network(monkeypatch, exc = None): + """Fail every urlopen and return a counter; tests assert n == 0 to prove no fetch + happened (check_upgrade_for_model swallows exceptions, so a raising stub alone + cannot prove the negative).""" + calls = {"n": 0} + + def _raise(*args, **kwargs): + calls["n"] += 1 + raise (exc or OSError("network fetch attempted")) + + monkeypatch.setattr("urllib.request.urlopen", _raise) + return calls + + +# --- AST extraction shared with the static router --- + + +class TestModelTypesFromSource: + def test_ordereddict_update_and_unpacking(self): + keys = _model_types_from_source(_MAPPING_SOURCE) + assert keys == {"llama", "gemma4", "qwen3_moe", "brandnew_arch"} + + def test_plain_dict_literal(self): + keys = _model_types_from_source(_MAIN_ONLY_SOURCE) + assert "dev_only_arch" in keys and "llama" in keys + + def test_syntax_error_raises_for_caller_to_handle(self): + with pytest.raises(SyntaxError): + _model_types_from_source("def broken(:\n") + + +class TestFetchRemoteModelTypes: + def test_merges_both_auto_files(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + keys = _fetch_remote_model_types("v5.13.0") + assert keys is not None and "brandnew_arch" in keys + + def test_all_fetches_failing_returns_none(self, monkeypatch): + _no_network(monkeypatch, exc = OSError("no route")) + assert _fetch_remote_model_types("main") is None + + def test_empty_mapping_treated_as_failure(self, monkeypatch): + monkeypatch.setattr( + "urllib.request.urlopen", + lambda req, timeout = None: _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n"), + ) + assert _fetch_remote_model_types("main") is None + + def test_transient_failure_of_one_file_fails_whole_lookup(self, monkeypatch): + # One file times out: the partial map must not be returned and cached. + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + raise OSError("timed out") + + monkeypatch.setattr("urllib.request.urlopen", _fake) + assert _fetch_remote_model_types("main") is None + + def test_missing_auto_mappings_404_still_succeeds(self, monkeypatch): + # Pre-5.10 tags have no auto_mappings.py; a 404 must not fail the lookup. + import urllib.error + + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + raise urllib.error.HTTPError(url, 404, "Not Found", None, None) + + monkeypatch.setattr("urllib.request.urlopen", _fake) + keys = _fetch_remote_model_types("v5.9.0") + assert keys is not None and "brandnew_arch" in keys + + def test_unparseable_file_fails_whole_lookup(self, monkeypatch): + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + return _FakeResponse(b"def broken(:\n") + + monkeypatch.setattr("urllib.request.urlopen", _fake) + assert _fetch_remote_model_types("main") is None + + +# --- latest_transformers_supports: snapshot, cache, offline, kill switch --- + + +class TestLatestTransformersSupports: + def test_supported_in_pypi(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("brandnew_arch") + assert result == { + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + } + + def test_dev_only_arch_reported_main_only(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("dev_only_arch") + assert result["supported_in_pypi"] is False + assert result["supported_in_main"] is True + + def test_unknown_everywhere(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("no_such_arch") + assert result["supported_in_pypi"] is False and result["supported_in_main"] is False + + def test_network_failure_returns_none(self, monkeypatch): + _no_network(monkeypatch, exc = OSError("down")) + assert latest_transformers_supports("brandnew_arch") is None + + def test_offline_returns_none_without_fetch(self, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + calls = _no_network(monkeypatch) + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == 0 + + def test_kill_switch_returns_none_without_fetch(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + calls = _no_network(monkeypatch) + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == 0 + + def test_memory_cache_hit_avoids_refetch(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + first_total = counter["__total__"] + latest_transformers_supports("some_other_arch") + assert counter["__total__"] == first_total + + def test_disk_cache_survives_restart(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + # Simulate a restart: memory gone, disk snapshot stays, network unavailable. + tl.clear_caches() + _no_network(monkeypatch) + result = latest_transformers_supports("brandnew_arch") + assert result is not None and result["supported_in_pypi"] is True + + def test_expired_snapshot_refetches(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + stale = dict(tl._memory_snapshot, fetched_at = time.time() - tl._CACHE_TTL_SECONDS - 1) + tl.clear_caches() + tl._save_snapshot_file(stale) + first_total = counter["__total__"] + latest_transformers_supports("brandnew_arch") + assert counter["__total__"] > first_total + + def test_corrupt_disk_cache_ignored(self, monkeypatch, tmp_path: Path): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + tl._cache_file().write_text("{not json", encoding = "utf-8") + result = latest_transformers_supports("brandnew_arch") + assert result is not None and counter["__total__"] > 0 + + def test_failure_backoff_skips_immediate_retry(self, monkeypatch): + calls = {"n": 0} + + def _fail(*args, **kwargs): + calls["n"] += 1 + raise OSError("down") + + monkeypatch.setattr("urllib.request.urlopen", _fail) + assert latest_transformers_supports("brandnew_arch") is None + first = calls["n"] + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == first # backed off, no second network attempt + + +# --- check_upgrade_for_model: the tier hook --- + + +def _local_model(tmp_path: Path, model_type: str) -> str: + d = tmp_path / f"model_{model_type}" + d.mkdir() + (d / "config.json").write_text(json.dumps({"model_type": model_type})) + return str(d) + + +_FAKE_OVERLAYS = { + "default": frozenset({"llama", "bert", "gpt2"}), + "530": frozenset({"qwen3_moe", "qwen3_next"}), + "550": frozenset({"gemma4"}), + "510": frozenset({"gemma4_unified"}), + "latest": frozenset(), +} + + +def _fake_overlays(monkeypatch, overlays = None): + overlays = overlays or _FAKE_OVERLAYS + fake = lambda tier: overlays.get(tier, frozenset()) + monkeypatch.setattr(tv, "_config_model_types", fake) + monkeypatch.setattr(tl, "_config_model_types", fake) + + +class TestCheckUpgradeForModel: + def test_unknown_type_supported_in_pypi_signals(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) + assert result == { + "model_type": "brandnew_arch", + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + } + + def test_dev_only_type_signals_main_only(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = check_upgrade_for_model(_local_model(tmp_path, "dev_only_arch")) + assert result["supported_in_pypi"] is False and result["supported_in_main"] is True + + def test_unknown_everywhere_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + assert check_upgrade_for_model(_local_model(tmp_path, "no_such_arch")) is None + + def test_offline_falls_through_without_fetch(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + assert calls["n"] == 0 + + def test_network_failure_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch, exc = OSError("down")) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + + def test_known_default_type_never_fetches(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "llama")) is None + assert calls["n"] == 0 + + def test_known_sidecar_type_never_fetches(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "gemma4_unified")) is None + assert calls["n"] == 0 + + def test_hardcoded_tier_type_never_fetches_even_without_overlays( + self, tmp_path: Path, monkeypatch + ): + # Sidecar overlays unreadable, but the hardcoded tables route it. + _fake_overlays( + monkeypatch, + {"default": frozenset({"llama"})}, + ) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "qwen3_5_moe")) is None + assert calls["n"] == 0 + + def test_unreadable_default_overlay_bails_out(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch, {"default": frozenset()}) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + assert calls["n"] == 0 + + def test_no_model_type_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch) + d = tmp_path / "no_type" + d.mkdir() + (d / "config.json").write_text(json.dumps({"architectures": ["Whatever"]})) + assert check_upgrade_for_model(str(d)) is None + + def test_nested_model_type_is_used(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + d = tmp_path / "nested" + d.mkdir() + (d / "config.json").write_text(json.dumps({"text_config": {"model_type": "brandnew_arch"}})) + result = check_upgrade_for_model(str(d)) + assert result is not None and result["model_type"] == "brandnew_arch" + + def test_never_raises_on_internal_error(self, monkeypatch): + monkeypatch.setattr( + tl, "_load_config_json", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) + ) + assert check_upgrade_for_model("some/model") is None + + +class TestNestedModelTypeExtraction: + def test_top_level_wins(self): + assert _model_types_from_config( + {"model_type": "a", "text_config": {"model_type": "b"}} + ) == ["a", "b"] + + def test_nested_fallback(self): + assert _model_types_from_config({"llm_config": {"model_type": "b"}}) == ["b"] + + def test_missing_returns_none(self): + assert _model_types_from_config({}) == [] + + +# --- Routing parity: overlay-shipped model_types route as before, never remote-check --- + + +class TestRoutingParity: + def test_all_overlay_types_route_identically_and_never_check(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + expected_tier = { + "llama": "default", + "bert": "default", + "gpt2": "default", + "qwen3_moe": "530", + "qwen3_next": "530", + "gemma4": "550", + "gemma4_unified": "510", + } + for model_type, tier in expected_tier.items(): + cfg = {"model_type": model_type} + assert _tier_from_config_mapping(cfg) == tier, model_type + assert check_upgrade_for_model(_local_model(tmp_path, model_type)) is None + assert calls["n"] == 0 + + def test_real_installed_mappings_route_without_checker(self, monkeypatch, tmp_path: Path): + """Parity over the REAL installed overlays (base + any provisioned sidecar): + every shipped model_type resolves statically, so the remote checker never + fires and routing is byte-identical with the feature enabled.""" + _no_network(monkeypatch) + seen = 0 + for tier in ("default", "530", "550", "510"): + types = tv._config_model_types(tier) + if not types: + continue # overlay not provisioned in this environment + for model_type in types: + assert _tier_from_config_mapping({"model_type": model_type}) is not None + seen += 1 + if seen == 0: + pytest.skip("no transformers overlay available in this environment") + + def test_get_tier_unchanged_by_kill_switch(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch) + path = _local_model(tmp_path, "no_such_arch") + _config_json_cache.clear() + tier_default = get_transformers_tier(path, probe = False) + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + _config_json_cache.clear() + assert get_transformers_tier(path, probe = False) == tier_default == "default" + + +# --- .venv_t5_latest provisioning and routing participation --- + + +class TestLatestVenvProvisioning: + def test_version_string_validation(self): + assert _is_valid_version_string("5.13.0") + assert _is_valid_version_string("5.14.0rc1") + assert not _is_valid_version_string("5.13.0; rm -rf /") + assert not _is_valid_version_string("git+https://evil") + assert not _is_valid_version_string("") + + def test_packages_pin_exact_version(self): + pkgs = _venv_t5_latest_packages("5.13.0") + assert pkgs[0] == "transformers==5.13.0" + assert any(p.startswith("huggingface_hub==") for p in pkgs) + + def test_ensure_latest_writes_pin_and_invalidates_cache(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + recorded = {} + + def _fake_ensure(dir_, packages, label): + recorded["dir"] = dir_ + recorded["packages"] = packages + Path(dir_).mkdir(parents = True, exist_ok = True) + return True + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + _config_mapping_cache["latest"] = frozenset({"stale"}) + assert ensure_latest_transformers_venv("5.13.0") is True + # Stage-and-swap: pip installs into staging, the live dir is the swap result. + assert recorded["dir"] == str(venv_dir) + ".staging" + assert "transformers==5.13.0" in recorded["packages"] + assert venv_dir.is_dir() + assert not Path(str(venv_dir) + ".staging").exists() + assert latest_venv_pinned_version() == "5.13.0" + assert "latest" not in _config_mapping_cache + + def test_ensure_latest_upgrade_failure_keeps_old_sidecar(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + venv_dir.mkdir(parents = True) + (venv_dir / tv._LATEST_PIN_MARKER).write_text( + json.dumps({"version": "5.12.0", "packages": ["transformers==5.12.0"]}) + ) + (venv_dir / "transformers").mkdir() + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a, **k: True) + # Install fails mid-flight: the previous sidecar and pin survive. + monkeypatch.setattr(tv, "_ensure_venv_dir", lambda *a, **k: False) + assert ensure_latest_transformers_venv("5.13.0") is False + assert latest_venv_pinned_version() == "5.12.0" + assert (venv_dir / "transformers").is_dir() + assert not Path(str(venv_dir) + ".staging").exists() + + def test_ensure_latest_rejects_bad_version(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert ensure_latest_transformers_venv("5.13.0 && curl evil") is False + + def test_ensure_latest_offline_refuses(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert ensure_latest_transformers_venv("5.13.0") is False + + def test_unpinned_sidecar_never_installs(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert tv._ensure_venv_t5_latest_exists() is False + + def test_pinned_sidecar_repairs_with_same_version(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a: False) + recorded = {} + + def _fake_ensure(dir_, packages, label): + recorded["dir"] = dir_ + recorded["packages"] = packages + Path(dir_).mkdir(parents = True, exist_ok = True) + return True + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + assert tv._ensure_venv_t5_latest_exists() is True + # Repair also stage-and-swaps, never installing into the live dir. + assert recorded["dir"] == str(venv_dir) + ".staging" + assert "transformers==5.13.0" in recorded["packages"] + assert latest_venv_pinned_version() == "5.13.0" + + +class TestLatestTierRouting: + def test_latest_outranks_510(self): + assert _higher_tier("latest", "510") == "latest" + assert _higher_tier("510", "latest") == "latest" + + def test_tier_from_mapping_prefers_lowest_but_reaches_latest(self, monkeypatch): + overlays = dict(_FAKE_OVERLAYS) + overlays["latest"] = frozenset({"brandnew_arch"}) + _fake_overlays(monkeypatch, overlays) + assert _tier_from_config_mapping({"model_type": "brandnew_arch"}) == "latest" + # Anything a lower tier ships stays on the lower tier. + assert _tier_from_config_mapping({"model_type": "qwen3_moe"}) == "530" + + def test_overlay_dir_for_latest(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + (venv_dir / "transformers").mkdir(parents = True) + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + # Unpinned dir is ignored: activation refuses an unpinned sidecar. + assert tv._overlay_transformers_dir("latest") is None + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") + + def test_probe_order_excludes_unprovisioned_latest(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + + def test_probe_order_includes_provisioned_latest(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + ("latest",) + + def test_activation_prepends_latest_dir(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", lambda: True) + old_sys_path = list(sys.path) + old_pp = os.environ.get("PYTHONPATH") + try: + activate_transformers_for_subprocess("some/brand-new-model") + assert sys.path[0] == str(venv_dir) + assert os.environ["PYTHONPATH"].split(os.pathsep)[0] == str(venv_dir) + finally: + sys.path[:] = old_sys_path + if old_pp is None: + os.environ.pop("PYTHONPATH", None) + else: + os.environ["PYTHONPATH"] = old_pp + + def test_activation_raises_when_latest_missing(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + with pytest.raises(RuntimeError, match = "venv_t5_latest"): + activate_transformers_for_subprocess("some/brand-new-model") + + +# --- install_latest_transformers: the consent endpoint helper --- + + +class TestInstallLatestTransformers: + def test_success_path(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + recorded = {} + + def _fake_ensure( + version, + extra_packages = (), + before_swap = None, + ): + recorded["args"] = (version, extra_packages) + return True + + monkeypatch.setattr(tl, "ensure_latest_transformers_venv", _fake_ensure) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + result = install_latest_transformers("5.13.0") + assert result["success"] is True and result["version"] == "5.13.0" + assert recorded["args"] == ("5.13.0", ()) + + def test_version_mismatch_rejected(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + ) + result = install_latest_transformers("4.99.0") + assert result["success"] is False and "not the latest" in result["message"] + + def test_offline_rejected(self, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _no_network(monkeypatch) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "offline" in result["message"].lower() + + def test_kill_switch_rejected(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + _no_network(monkeypatch) + result = install_latest_transformers("5.13.0") + assert result["success"] is False + + def test_install_failure_reported(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (), before_swap = None: False, + ) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "failed" in result["message"] + + def test_blocked_by_incompatible_deps(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), ["numpy>=99.0"])) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + ) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "numpy>=99.0" in result["message"] + + def test_compat_shadows_passed_to_installer(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: (("tokenizers==0.23.0",), [])) + recorded = {} + + def _fake_ensure( + version, + extra_packages = (), + before_swap = None, + ): + recorded["extras"] = extra_packages + return True + + monkeypatch.setattr(tl, "ensure_latest_transformers_venv", _fake_ensure) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + result = install_latest_transformers("5.13.0") + assert result["success"] is True + assert recorded["extras"] == ("tokenizers==0.23.0",) + + +class TestCompatPlan: + def _patch_env(self, monkeypatch, requires, installed): + monkeypatch.setattr(tl, "_fetch_requires_dist", lambda v: requires) + + def _ver(name): + from importlib.metadata import PackageNotFoundError + + key = name.lower().replace("_", "-") + if key not in installed: + raise PackageNotFoundError(name) + return installed[key] + + monkeypatch.setattr("importlib.metadata.version", _ver) + + def test_satisfied_env_needs_nothing(self, monkeypatch): + self._patch_env( + monkeypatch, + ["tokenizers<=0.23.0,>=0.22.0", "safetensors>=0.8.0", "numpy>=1.17"], + {"tokenizers": "0.22.2", "safetensors": "0.8.0", "numpy": "2.4.4"}, + ) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + def test_unsatisfied_shadowable_dep_pinned(self, monkeypatch): + self._patch_env( + monkeypatch, + ["tokenizers>=0.24.0"], + {"tokenizers": "0.22.2"}, + ) + monkeypatch.setattr(tl, "_resolve_exact_version", lambda name, spec: "0.24.1") + extras, blockers = tl.compat_plan("5.99.0") + assert extras == ("tokenizers==0.24.1",) and blockers == [] + + def test_unsatisfied_non_shadowable_dep_blocks(self, monkeypatch): + self._patch_env(monkeypatch, ["numpy>=99.0"], {"numpy": "2.4.4"}) + extras, blockers = tl.compat_plan("5.99.0") + assert extras == () and blockers == ["numpy>=99.0"] + + def test_cli_only_dep_ignored(self, monkeypatch): + self._patch_env(monkeypatch, ["typer"], {}) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + def test_sidecar_provided_hub_checked_against_recipe_pin(self, monkeypatch): + self._patch_env(monkeypatch, ["huggingface-hub<2.0,>=1.5.0"], {"huggingface-hub": "0.36.2"}) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] # 1.8.0 sidecar pin satisfies it + + def test_sidecar_provided_hub_out_of_range_blocks(self, monkeypatch): + self._patch_env(monkeypatch, ["huggingface-hub>=2.1"], {"huggingface-hub": "0.36.2"}) + extras, blockers = tl.compat_plan("5.99.0") + assert blockers == ["huggingface-hub>=2.1"] + + def test_unfetchable_requires_dist_blocks_install(self, monkeypatch): + # Proceeding unverified could pin a sidecar whose imports crash workers. + monkeypatch.setattr(tl, "_fetch_requires_dist", lambda v: None) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and len(blockers) == 1 and "retry" in blockers[0] + + def test_extra_marker_requirements_skipped(self, monkeypatch): + self._patch_env( + monkeypatch, + ['torch>=99.0; extra == "torch"', 'pytest; python_version < "3.0"'], + {}, + ) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + +def test_get_snapshot_dedupes_concurrent_fetch(monkeypatch): + """While one thread is fetching, other callers return None instead of stacking fetches.""" + with tl._lock: + tl._is_fetching = True + calls = {"n": 0} + + def boom(): + calls["n"] += 1 + raise AssertionError("must not fetch while another fetch is in flight") + + monkeypatch.setattr(tl, "_refresh_snapshot", boom) + assert tl._get_snapshot() is None + assert calls["n"] == 0 + tl.clear_caches() + + +def test_install_serialized(): + """A second install call while one is in progress gets a structured refusal.""" + from utils.transformers_version import try_begin_sidecar_swap + + assert try_begin_sidecar_swap() is True + out = tl.install_latest_transformers("5.13.0") + assert out["success"] is False + assert "already in progress" in out["message"] + tl.clear_caches() + + +def test_install_in_progress_reflects_reservation(): + """is_install_in_progress mirrors the shared sidecar swap reservation, so a + lazy repair (which takes the same reservation) also blocks worker starts.""" + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + assert tl.is_install_in_progress() is False + assert try_begin_sidecar_swap() is True + try: + assert tl.is_install_in_progress() is True + finally: + end_sidecar_swap() + assert tl.is_install_in_progress() is False + + +def test_upgrade_check_sees_nested_model_types(monkeypatch): + """A supported wrapper with a brand-new nested backbone must still signal.""" + cfg = { + "model_type": "llava", # in every installed overlay + "text_config": {"model_type": "zz_brand_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_brand_new_llm", + "supported_in_main": mt == "zz_brand_new_llm", + }, + ) + out = tl.check_upgrade_for_model("some-org/wrapped-new-backbone") + assert out is not None + assert out["model_type"] == "zz_brand_new_llm" + + +def test_upgrade_check_ignores_nested_known_types(monkeypatch): + """All nested types known to installed overlays -> no signal, no remote call.""" + cfg = { + "model_type": "llava", + "text_config": {"model_type": "llama"}, + "vision_config": {"model_type": "clip_vision_model"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + calls = [] + monkeypatch.setattr(tl, "latest_transformers_supports", lambda mt: calls.append(mt) or None) + assert tl.check_upgrade_for_model("some-org/normal-vlm") is None + assert calls == [] + + +def test_upgrade_check_requires_primary_supported(monkeypatch): + """Latest supporting only a nested type must not prompt: routing still + cannot load the primary, so the install would not fix the model.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_llm", + "supported_in_main": mt == "zz_new_llm", + }, + ) + assert tl.check_upgrade_for_model("some-org/half-supported") is None + + +def test_upgrade_check_requires_every_missing_type(monkeypatch): + """Primary supported but a nested backbone missing from latest -> no prompt + (CONFIG_MAPPING would still fail on the sub-config); all supported -> signal + carries the primary type.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_wrapper", + "supported_in_main": mt == "zz_new_wrapper", + }, + ) + assert tl.check_upgrade_for_model("some-org/half-supported") is None + + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + }, + ) + out = tl.check_upgrade_for_model("some-org/fully-supported") + assert out is not None and out["model_type"] == "zz_new_wrapper" + + +def test_install_success_invalidates_capability_caches(monkeypatch): + """A successful install must drop tier probes, the latest mapping, and the + vision-detection cache so the new sidecar takes effect without a restart.""" + from utils.models import model_config as mc + + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + monkeypatch.setattr( + tl, "ensure_latest_transformers_venv", lambda v, extra_packages = (), before_swap = None: True + ) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + + tv._probe_tier_cache["stale/model"] = "default" + tv._config_mapping_cache["latest"] = frozenset({"stale_type"}) + tv._config_mapping_cache["default"] = frozenset({"llama"}) + mc._vision_detection_cache[("stale/model", None, False)] = False + + result = install_latest_transformers("5.13.0") + assert result["success"] is True + assert tv._probe_tier_cache == {} + assert "latest" not in tv._config_mapping_cache + assert tv._config_mapping_cache.get("default") == frozenset({"llama"}) # untouched + assert mc._vision_detection_cache == {} + + tv._probe_tier_cache.clear() + tv._config_mapping_cache.clear() + tl.clear_caches() + + +def test_vision_subprocess_unions_sidecar_registry(): + """The embedded vision-check script must extend the inlined parent sets with + the ACTIVE sidecar's registry so sidecar-only architectures classify.""" + from utils.models import model_config as mc + + script = mc._VISION_CHECK_SCRIPT + ast.parse(script) + stub_registry = { + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES": { + "zz_sidecar_vlm": "ZzSidecarForConditionalGeneration" + }, + } + ns = {} + # Exec only the registry-union block against a stubbed sidecar registry. + body = script.split("from transformers import AutoConfig", 1)[1] + body = body.split("kwargs = {", 1)[0] + helpers = script.split("sys.path.insert(0, backend_dir)", 1)[1] + helpers = helpers.split("try:", 1)[0] + exec(helpers, ns) + + class _FakeMa: + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = stub_registry[ + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES" + ] + + import sys as _sys + import types as _types + + fake_pkg = _types.ModuleType("transformers.models.auto") + fake_pkg.modeling_auto = _FakeMa + saved = { + k: _sys.modules.get(k) + for k in ("transformers.models.auto", "transformers.models.auto.modeling_auto") + } + _sys.modules["transformers.models.auto"] = fake_pkg + _sys.modules["transformers.models.auto.modeling_auto"] = _FakeMa + try: + exec(textwrap.dedent(body), ns) + finally: + for k, v in saved.items(): + if v is None: + _sys.modules.pop(k, None) + else: + _sys.modules[k] = v + + assert "zz_sidecar_vlm" in ns["_VLM_MODEL_TYPES"] + assert "ZzSidecarForConditionalGeneration" in ns["_VLM_CLASS_NAMES"] + + class _Cfg: + architectures = ["ZzSidecarForConditionalGeneration"] + model_type = "zz_sidecar_vlm" + + assert ns["_is_vlm"](_Cfg()) is True + + +def test_upgrade_check_mixed_pypi_main_reports_dev_only(monkeypatch): + """Primary in the PyPI release but a nested type only on main: no install + may be offered (CONFIG_MAPPING would fail on the nested sub-config), so the + aggregate must read as main-only.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_wrapper", + "supported_in_main": True, + }, + ) + out = tl.check_upgrade_for_model("some-org/mixed-support") + assert out is not None + assert out["model_type"] == "zz_new_wrapper" + assert out["supported_in_pypi"] is False # no install offered + assert out["supported_in_main"] is True + + +def test_install_endpoint_not_mounted_on_v1(): + """The consented pip-install endpoint is a Studio admin action; it must live + on studio_router (kept off the OpenAI-compatible /v1 mount), not router.""" + from routes import inference as ri + + path = "/install-latest-transformers" + assert path in [r.path for r in ri.studio_router.routes] + assert path not in [r.path for r in ri.router.routes] + + +def test_kill_switch_removes_provisioned_latest_from_routing(tmp_path, monkeypatch): + """UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS must roll back a provisioned latest + sidecar: no overlay mapping, no probe participation, no file deletion needed.""" + venv_dir = tmp_path / ".venv_t5_latest" + (venv_dir / "transformers").mkdir(parents = True) + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + + assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + ("latest",) + + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + tv._config_mapping_cache.pop("latest", None) + assert tv._overlay_transformers_dir("latest") is None + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + tv._config_mapping_cache.pop("latest", None) + + +def test_repair_failure_preserves_pin_and_live_dir(tmp_path, monkeypatch): + """A failed lazy repair must not delete the incomplete-but-pinned live + sidecar: the pin survives so a later attempt can still repair it.""" + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + (venv_dir / "partial_file").write_text("x") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a: False) + monkeypatch.setattr(tv, "_ensure_venv_dir", lambda *a, **k: False) + + from utils.transformers_version import latest_venv_pinned_version + + assert tv._ensure_venv_t5_latest_exists() is False + assert venv_dir.is_dir() + assert (venv_dir / "partial_file").exists() + assert latest_venv_pinned_version() == "5.13.0" + assert not (tmp_path / ".venv_t5_latest.staging").exists() + + +def test_failed_staging_install_removes_staging_dir(tmp_path, monkeypatch): + """A pip failure inside _ensure_venv_dir returns False without raising, so + the except cleanup never runs; the partial staging dir must still go.""" + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + + def _fake_ensure(dir_, packages, label): + Path(dir_).mkdir(parents = True, exist_ok = True) + (Path(dir_) / "partial").write_text("x") + return False + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + assert ensure_latest_transformers_venv("5.13.0") is False + assert not Path(str(venv_dir) + ".staging").exists() diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index b9b5abb9e5..a6e6803a5c 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -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" diff --git a/studio/backend/tests/test_windows_external_drive_paths.py b/studio/backend/tests/test_windows_external_drive_paths.py new file mode 100644 index 0000000000..9686d45c9f --- /dev/null +++ b/studio/backend/tests/test_windows_external_drive_paths.py @@ -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, "", "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, "", "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, "", "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 diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 281ca24281..284bbb5745 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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 diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py index 1f1754664f..0ea0477cc7 100644 --- a/studio/backend/utils/paths/external_media.py +++ b/studio/backend/utils/paths/external_media.py @@ -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] diff --git a/studio/backend/utils/transformers_latest.py b/studio/backend/utils/transformers_latest.py new file mode 100644 index 0000000000..40c8f729a5 --- /dev/null +++ b/studio/backend/utils/transformers_latest.py @@ -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 diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index a69673f081..9f9f8aa3de 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -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 diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 6f68917224..ba56ce7525 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -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() { {!isAuthFlowRoute && } + {hideNavbar ? (
}> diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 5b6264c6d7..9b502a5000 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -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} > - {/* Active-mode badge: always visible when bypass is on, even while - the pill row is collapsed (returns null when off). */} - + {/* Permission-level pill: always visible, even while the pill row + is collapsed; opens the permission level dropdown. */} + {composerExpanded ? ( <> @@ -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 ( - - ); -}; - const ToolStatusDisplay: FC = () => { const toolStatus = useChatRuntimeStore((s) => s.toolStatus); const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 9ed11deccc..c7dd6372aa 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -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). diff --git a/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx index 14cb6747e9..b35317b2fa 100644 --- a/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx +++ b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx @@ -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 ( - { - 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); + + - - Bypass permissions - {bypassPermissions ? ( - - ) : null} - + > + + Bypass permissions + + + + setTimeout(() => setBypassConfirmOpen(true), 0) + } + /> + + ); } @@ -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 ( - Enable Bypass permissions? + Enable 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 + 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 @@ -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); }} > diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 07ddffdd59..cedd298ecf 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -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 (
@@ -2049,85 +2039,49 @@ function ConfirmToolCallsToggle() { Confirm tool calls - 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.
- {bypassPermissions ? ( + {permissionMode === "full" ? ( - Overridden by Bypass permissions + Overridden by Full access (Bypass permissions) ) : null} ); } 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 ( -
-
-
- - Bypass permissions - - - 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. - -
- { - if (next) setDialogOpen(true); - else setBypassPermissions(false); - }} - /> +
+
+ + Bypass permissions + + + How Unsloth approves tool calls before they run. Full access is + dangerous: it disables confirmations and the code sandbox. +
- {bypassPermissions ? ( + {/* Full width, styled like the panel selects/preset input. */} + + {permissionMode === "full" ? ( Tool calls run with no confirmation and no sandbox. ) : null} - - - - Enable Bypass permissions? - - 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 - - - - Cancel - { - setBypassPermissions(true); - setDialogOpen(false); - }} - > - I understand - - - -
); } diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index a659b7f83e..4646f67cb4 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -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. diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 3099884645..7e894bb92e 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -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"; diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx new file mode 100644 index 0000000000..4277c1bfcf --- /dev/null +++ b/studio/frontend/src/features/chat/permission-mode-select.tsx @@ -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) => ( + { + // 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.label} + + {option.description} + + + {permissionMode === option.value ? ( + + ) : null} + + ))} + + ); +} + +/** 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 ( + + + + Enable Full access? + + 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 + + + + Cancel + { + setPermissionMode("full"); + onOpenChange(false); + }} + > + I understand + + + + + ); +} + +/** + * 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 ( + <> + + + + + + + How should tool calls be approved? + + + setTimeout(() => setConfirmOpen(true), 0) + } + /> + + + + + ); +} + +/** + * 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 ( + + + + + + + How should tool calls be approved? + + + setTimeout(() => setBypassConfirmOpen(true), 0) + } + /> + + + ); +} diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index a50af46e85..2ed9589461 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -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({ Compare - {/* 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 && ( - - )} + {/* 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. */} +